src/utils.py

"""
Utility functions for Nova AI - ChatGPT Astra Edition.
Hardware inspection, Astra prompt enhancement, latency profiling, and chat exporters (MD, JSON, HTML).
"""

import os
import sys
import time
import json
import html
import psutil
from pathlib import Path
from typing import List, Dict, Any, Tuple

# Ensure UTF-8 output on Windows consoles
if sys.platform == "win32":
    try:
        if sys.stdout and hasattr(sys.stdout, "reconfigure"):
            sys.stdout.reconfigure(encoding="utf-8")
        if sys.stderr and hasattr(sys.stderr, "reconfigure"):
            sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

from rich.console import Console

console = Console(legacy_windows=False)


def get_system_hardware_info() -> Dict[str, Any]:
    """Inspect and return current CPU, RAM, NVIDIA GPU, and Hugging Face ZeroGPU capabilities."""
    info: Dict[str, Any] = {
        "cpu_count_physical": psutil.cpu_count(logical=False) or 1,
        "cpu_count_logical": psutil.cpu_count(logical=True) or 1,
        "total_ram_gb": round(psutil.virtual_memory().total / (1024 ** 3), 2),
        "available_ram_gb": round(psutil.virtual_memory().available / (1024 ** 3), 2),
        "has_cuda": False,
        "is_zero_gpu": False,
        "device_name": "CPU (Intel / AMD Multi-Core)",
        "recommended_device": "cpu",
        "cuda_device_count": 0,
        "cuda_vram_gb": 0.0,
        "cuda_version": None,
        "compute_capability": None,
        "hardware_badge": "⚡ Multi-Core CPU Accelerated"
    }

    # Check for Hugging Face Spaces ZeroGPU environment
    if os.getenv("SPACES_ZERO_GPU") == "true" or "spaces" in sys.modules:
        info["is_zero_gpu"] = True
        info["hardware_badge"] = "🟢 NVIDIA A100 ZeroGPU Dynamic Allocation Active"

    try:
        import torch
        if torch.cuda.is_available():
            info["has_cuda"] = True
            info["cuda_device_count"] = torch.cuda.device_count()
            dev_name = torch.cuda.get_device_name(0)
            info["device_name"] = dev_name
            info["recommended_device"] = "cuda"
            vram = round(torch.cuda.get_device_properties(0).total_memory / (1024 ** 3), 2)
            info["cuda_vram_gb"] = vram
            info["cuda_version"] = torch.version.cuda
            cap = torch.cuda.get_device_capability(0)
            info["compute_capability"] = f"{cap[0]}.{cap[1]}"

            if info["is_zero_gpu"]:
                info["hardware_badge"] = f"🟢 NVIDIA ZeroGPU ({dev_name}) Active"
            else:
                info["hardware_badge"] = f"🟢 NVIDIA CUDA ({dev_name} • {vram}GB VRAM • CC {cap[0]}.{cap[1]})"

        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            info["device_name"] = "Apple Silicon MPS Metal"
            info["recommended_device"] = "mps"
            info["hardware_badge"] = "🍎 Apple Silicon Metal MPS Accelerated"
    except ImportError:
        pass

    return info


def enhance_prompt_with_astra(prompt: str) -> str:
    """
    Astra Spark Prompt Enhancer:
    Transforms a concise prompt into a structured, high-yield ChatGPT Astra prompt.
    """
    cleaned = prompt.strip()
    if not cleaned:
        return prompt

    # If it's already a very detailed prompt, don't over-wrap it
    if len(cleaned) > 250 and ("###" in cleaned or "Requirements:" in cleaned):
        return cleaned

    # Generate an elevated Astra prompt template
    enhanced = (
        f"{cleaned}\n\n"
        f"--- Astra Guidance ---\n"
        f"Please provide an in-depth, structured response covering:\n"
        f"1. Core concept and fundamental explanation.\n"
        f"2. Practical real-world examples or code implementation with best practices.\n"
        f"3. Key trade-offs, common pitfalls, and edge cases to consider.\n"
        f"4. Summary checklist or actionable next steps."
    )
    return enhanced


def format_chat_as_markdown(history: List[Dict[str, str]], title: str = "Nova AI - Conversation Record") -> str:
    """Format conversation list into structured Markdown text."""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    md = [
        f"# 🌟 {title}\n",
        f"*Generated by Nova AI (ChatGPT Astra Edition) on: {timestamp}*\n\n",
        "---\n\n"
    ]
    for msg in history:
        role = msg.get("role", "user").lower()
        content = msg.get("content", "").strip()
        if role == "system":
            md.append(f"> 🧠 **Astra System Persona:**\n> {content}\n\n")
        elif role == "user":
            md.append(f"### 👤 User:\n{content}\n\n")
        else:
            md.append(f"### 🌟 Nova Astra:\n{content}\n\n")
    return "".join(md)


def format_chat_as_json(history: List[Dict[str, str]]) -> str:
    """Format conversation history into rich JSON string."""
    return json.dumps({
        "generator": "Nova AI - ChatGPT Astra Edition",
        "version": "2.5.0-Astra",
        "timestamp": time.time(),
        "date": time.strftime("%Y-%m-%d %H:%M:%S"),
        "total_messages": len(history),
        "messages": history
    }, indent=2, ensure_ascii=False)


def format_chat_as_html(history: List[Dict[str, str]], title: str = "Nova AI Conversation") -> str:
    """Format conversation into an exportable standalone cosmic dark-mode HTML document."""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    messages_html = []

    for msg in history:
        role = msg.get("role", "user").lower()
        content = html.escape(msg.get("content", "").strip()).replace("\n", "<br/>")

        if role == "system":
            card = f"""
            <div class="msg-box system">
                <div class="msg-header">🧠 Astra Persona Directive</div>
                <div class="msg-body">{content}</div>
            </div>
            """
        elif role == "user":
            card = f"""
            <div class="msg-box user">
                <div class="msg-header">👤 User</div>
                <div class="msg-body">{content}</div>
            </div>
            """
        else:
            card = f"""
            <div class="msg-box astra">
                <div class="msg-header">🌟 Nova Astra</div>
                <div class="msg-body">{content}</div>
            </div>
            """
        messages_html.append(card)

    html_content = f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>{html.escape(title)} - Nova AI</title>
    <style>
        :root {{
            --bg-color: #0b0f19;
            --card-bg: #111827;
            --text-main: #f3f4f6;
            --text-dim: #9ca3af;
            --accent-purple: #8b5cf6;
            --accent-cyan: #06b6d4;
            --border-color: #1f2937;
        }}
        body {{
            background-color: var(--bg-color);
            color: var(--text-main);
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            line-height: 1.6;
            margin: 0;
            padding: 30px 15px;
        }}
        .container {{
            max-width: 900px;
            margin: 0 auto;
        }}
        .header {{
            background: linear-gradient(135deg, #1e1b4b 0%, #312e81 50%, #4338ca 100%);
            padding: 24px 30px;
            border-radius: 16px;
            margin-bottom: 24px;
            box-shadow: 0 10px 25px -5px rgba(49, 46, 129, 0.4);
        }}
        .header h1 {{
            margin: 0 0 6px 0;
            font-size: 26px;
            color: #fff;
        }}
        .header p {{
            margin: 0;
            color: #c7d2fe;
            font-size: 14px;
        }}
        .msg-box {{
            margin-bottom: 16px;
            padding: 16px 20px;
            border-radius: 14px;
            background: var(--card-bg);
            border: 1px solid var(--border-color);
        }}
        .msg-box.user {{
            border-left: 4px solid var(--accent-cyan);
        }}
        .msg-box.astra {{
            border-left: 4px solid var(--accent-purple);
            background: #141c2e;
        }}
        .msg-box.system {{
            border-left: 4px solid #f59e0b;
            font-size: 13px;
            color: var(--text-dim);
        }}
        .msg-header {{
            font-weight: 700;
            font-size: 13px;
            text-transform: uppercase;
            letter-spacing: 0.05em;
            margin-bottom: 8px;
            color: #e0e7ff;
        }}
        .msg-body {{
            white-space: pre-wrap;
            word-break: break-word;
        }}
        .footer {{
            text-align: center;
            color: var(--text-dim);
            font-size: 13px;
            margin-top: 30px;
        }}
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>🌟 Nova AI (ChatGPT Astra Edition)</h1>
            <p>Exported Transcript • {timestamp}</p>
        </div>
        {''.join(messages_html)}
        <div class="footer">Generated by Nova AI • Next-Gen Conversational Intelligence</div>
    </div>
</body>
</html>"""
    return html_content


class LatencyTracker:
    """High-resolution latency, time-to-first-token, and throughput tracker."""

    def __init__(self):
        self.start_time: float = 0.0
        self.first_token_time: float = 0.0
        self.end_time: float = 0.0
        self.token_count: int = 0

    def start(self):
        self.start_time = time.perf_counter()
        self.first_token_time = 0.0
        self.end_time = 0.0
        self.token_count = 0

    def record_token(self, count: int = 1):
        self.token_count += count
        if self.first_token_time == 0.0:
            self.first_token_time = time.perf_counter()

    def stop(self) -> Dict[str, float]:
        self.end_time = time.perf_counter()
        total_time = max(self.end_time - self.start_time, 0.0001)
        ttft = max(self.first_token_time - self.start_time, 0.0) if self.first_token_time > 0 else 0.0
        tps = self.token_count / total_time if self.token_count > 0 else 0.0

        return {
            "total_latency_seconds": round(total_time, 3),
            "time_to_first_token_ms": round(ttft * 1000, 1),
            "tokens_generated": self.token_count,
            "tokens_per_second": round(tps, 1)
        }