"""
Interactive CLI Chat Interface for Nova AI - ChatGPT Astra Edition.
Rich-powered terminal interface featuring live token streaming, Astra Deep Reasoning,
token vault management, multimodal generation (/image, /web), and session controls.
"""

import os
import sys
from pathlib import Path
from typing import List, Dict, Any

# 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

# Ensure project root in sys.path
PROJECT_ROOT = Path(__file__).resolve().parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.markdown import Markdown
from rich.prompt import Prompt

from src.config import config, AVAILABLE_MODELS, SYSTEM_PROMPT_PRESETS, GenerationConfig
from src.model_engine import engine
from src.token_system import token_vault, TOKEN_COSTS
from src.multimodal import generate_ai_image, generate_ai_website
from src.web_search import search_live_web
from src.utils import (
    format_chat_as_markdown,
    format_chat_as_json,
    format_chat_as_html,
    enhance_prompt_with_astra,
    get_system_hardware_info
)
from src.domain_registry import domain_registry

console = Console(legacy_windows=False)


def print_banner(reasoning_mode: bool = False, grounding_mode: bool = False, active_domain: str = None):
    active_model = engine.model_id or config.default_model_id
    reasoning_status = "[bold green]ENABLED[/bold green]" if reasoning_mode else "[dim]DISABLED[/dim]"
    grounding_status = "[bold green]ENABLED[/bold green]" if grounding_mode else "[dim]DISABLED[/dim]"
    domain_status = f"[bold magenta]{active_domain}[/bold magenta]" if active_domain else f"[bold green]{domain_registry.total_count:,} Domains[/bold green] [dim](Auto-Detect)[/dim]"
    balance = token_vault.get_balance()
    hw = get_system_hardware_info()
    banner_text = (
        "[bold cyan]🌟 NOVA AI — Flagship ChatGPT Astra Enterprise Suite[/bold cyan]\n"
        "[dim]Full 7B/8B Foundation Models • NVIDIA ZeroGPU & CUDA • Deep Reasoning • Multimodal Studio[/dim]\n\n"
        f"• Active Model: [bold yellow]{active_model}[/bold yellow]\n"
        f"• Hardware Engine: [bold green]{hw.get('hardware_badge')}[/bold green]\n"
        f"• Astra Tokens: [bold cyan]🪙 {balance} ASTRA[/bold cyan]\n"
        f"• Specialization: {domain_status}\n"
        f"• Deep Reasoning: {reasoning_status}\n"
        f"• Web Grounding: {grounding_status}\n\n"
        "[dim]Commands: [bold cyan]/help[/bold cyan], [bold yellow]/think[/bold yellow], [bold blue]/domains[/bold blue], [bold magenta]/image[/bold magenta], [bold blue]/web[/bold blue], [bold green]/tokens[/bold green], [bold red]/quit[/bold red][/dim]"
    )
    console.print(Panel(banner_text, border_style="cyan", title="[bold magenta]Astra Flagship Console[/bold magenta]"))


def print_help():
    help_md = """
### 🌟 Nova Astra CLI Commands:
- `/help`    - Display this command reference
- `/domains` - Browse and search 1,000+ domain specializations (1,425 available)
- `/domain`  - Lock into a specific domain specialization (e.g. `/domain upi` or `/domain auto`)
- `/train`   - Launch AI/ML Fine-Tuning Studio (PEFT LoRA + Weight Merging)
- `/tokens`  - View Astra token balance & transaction history
- `/think`   - Toggle Astra Deep Reasoning Mode (5 ASTRA)
- `/search`  - Live web search with DuckDuckGo & Wikipedia (3 ASTRA)
- `/ground`  - Toggle automatic real-time web grounding for chat
- `/image`   - Generate an AI Image with FLUX / SDXL (15 ASTRA)
- `/web`     - Generate a full responsive HTML5 website (20 ASTRA)
- `/spark`   - Enhance a prompt with Astra Spark Optimizer
- `/persona` - Switch ChatGPT Astra persona mode
- `/model`   - Switch active model architecture
- `/stats`   - View hardware profiler & session statistics
- `/export`  - Export dialogue transcript to Markdown, JSON, or HTML
- `/clear`   - Reset active conversation history
- `/quit`    - Exit interactive session
"""
    console.print(Markdown(help_md))


def cli_web_search():
    query = Prompt.ask("\n[bold cyan]🌐 Enter search query[/bold cyan]").strip()
    if not query:
        return
    with console.status("[bold cyan]Searching live web with DuckDuckGo & Wikipedia (3 ASTRA)...[/bold cyan]"):
        res = search_live_web(query)
    if res.get("status") == "success":
        results = res.get("results", [])
        console.print(f"[bold green]✓ Found {len(results)} live web sources:[/bold green]")
        for idx, r in enumerate(results, 1):
            console.print(f"[bold cyan][{idx}][/bold cyan] [bold white]{r['title']}[/bold white] [dim]({r['source']})[/dim]")
            console.print(f"    🔗 [link={r['url']}]{r['url']}[/link]")
            console.print(f"    📄 [dim]{r['snippet'][:180]}...[/dim]\n")
        console.print(f"🪙 Remaining Balance: [bold cyan]{res['remaining_tokens']} ASTRA[/bold cyan]")
    else:
        console.print(f"[bold red]Search error:[/bold red] {res.get('message')}")


def show_tokens():
    bal = token_vault.get_balance()
    history = token_vault.get_transaction_history()
    table = Table(title=f"🪙 Astra Token Vault (Balance: {bal} ASTRA)", border_style="cyan")
    table.add_column("Date", style="dim")
    table.add_column("Action", style="bold")
    table.add_column("Amount", justify="right")
    table.add_column("Balance After", justify="right", style="bold green")

    for t in history[:8]:
        amt_color = "green" if t.get("amount", 0) > 0 else "red"
        table.add_row(
            t.get("date", ""),
            t.get("action", ""),
            f"[{amt_color}]{t.get('amount', 0):+d}[/{amt_color}]",
            str(t.get("balance_after", ""))
        )
    console.print(table)


def cli_image_generator():
    prompt = Prompt.ask("\n[bold magenta]🎨 Enter image description[/bold magenta]").strip()
    if not prompt:
        return
    with console.status("[bold cyan]Synthesizing image with Astra FLUX neural engine (15 ASTRA)...[/bold cyan]"):
        res = generate_ai_image(prompt)
    if res.get("status") == "success":
        console.print(f"[bold green]✓ Image generated successfully![/bold green]")
        console.print(f"📁 Local File: [link={res['file_path']}]{res['file_path']}[/link]")
        console.print(f"🪙 Remaining Balance: [bold cyan]{res['remaining_tokens']} ASTRA[/bold cyan]")
    else:
        console.print(f"[bold red]Generation error:[/bold red] {res.get('message')}")


def cli_website_generator():
    prompt = Prompt.ask("\n[bold blue]🌐 Enter React & Next.js website concept[/bold blue]").strip()
    if not prompt:
        return
    with console.status("[bold cyan]Synthesizing React.js & Next.js project with Nova Astra WebCoder-V1 (20 ASTRA)...[/bold cyan]"):
        res = generate_ai_website(prompt)
    if res.get("status") == "success":
        console.print(f"[bold green]✓ React & Next.js Project generated successfully![/bold green]")
        console.print(f"📦 Auto ZIP Archive: [bold yellow]{res.get('zip_path')}[/bold yellow]")
        console.print(f"⚛️ React Component:  [bold cyan]{res.get('jsx_path')}[/bold cyan]")
        console.print(f"⚡ Next.js Page:     [bold magenta]{res.get('tsx_path')}[/bold magenta]")
        console.print(f"🌐 Live Sandbox:     [link={res['file_path']}]{res['file_path']}[/link]")
        console.print(f"🪙 Remaining Balance: [bold cyan]{res['remaining_tokens']} ASTRA[/bold cyan]")
    else:
        console.print(f"[bold red]Generation error:[/bold red] {res.get('message')}")


def cli_fine_tuning():
    from src.fine_tune import run_lora_finetuning
    from src.dataset_builder import compile_huge_training_dataset, HUGE_DATASET_FILE

    console.print(Panel(
        "[bold cyan]🧠 Nova Astra AI/ML Fine-Tuning Studio[/bold cyan]\n"
        "[dim]LoRA Parameter-Efficient Fine-Tuning on Huge Multi-Turn Datasets (React, Fintech INR, Multimodal, Reasoning)[/dim]",
        border_style="magenta"
    ))

    choice = Prompt.ask(
        "\n[bold yellow]Select Training Dataset Scope[/bold yellow]\n"
        "  [1] Quick LoRA Trial (10 samples, ~1-2 min)\n"
        "  [2] High-Yield Benchmark (50 samples, ~5 min)\n"
        "  [3] Massive Dataset (All 610 samples)\n"
        "  [4] Recompile / Verify Huge Dataset JSONL\n"
        "Choice",
        choices=["1", "2", "3", "4"],
        default="1"
    )

    if choice == "4":
        with console.status("[bold cyan]Compiling 600+ sample multi-turn dataset...[/bold cyan]"):
            d_res = compile_huge_training_dataset(target_count=600)
        console.print(f"[bold green]✓ Dataset compiled successfully:[/bold green] {d_res['sample_count']} samples (~{d_res['estimated_tokens']:,} tokens)")
        console.print(f"📁 Location: {d_res['file_path']}")
        return

    sample_map = {"1": 10, "2": 50, "3": None}
    max_samples = sample_map[choice]

    active_model = engine.model_id or config.default_model_id
    console.print(f"\n[bold cyan]Starting LoRA Fine-Tuning on {active_model} (samples: {max_samples or 'ALL'})...[/bold cyan]")
    with console.status("[bold green]Executing PEFT LoRA training loop & weight merging...[/bold green]"):
        res = run_lora_finetuning(
            base_model_id=active_model,
            epochs=1,
            batch_size=1,
            max_samples=max_samples,
            export_merged=True
        )

    if res.get("status") == "success":
        console.print("\n[bold green]🎉 Fine-Tuning & Weight Export Succeeded![/bold green]")
        console.print(f"  • Trained Samples: [bold white]{res.get('samples_trained')}[/bold white]")
        console.print(f"  • Training Loss: [bold cyan]{res.get('train_loss')}[/bold cyan]")
        console.print(f"  • Time Elapsed: [bold yellow]{res.get('training_time_seconds')}s[/bold yellow]")
        console.print(f"  • Adapter Dir: [dim]{res.get('adapter_dir')}[/dim]")
        console.print(f"  • Merged Safetensors: [bold green]{res.get('exported_model_dir')}[/bold green]")
    else:
        console.print(f"[bold red]❌ Training failed:[/bold red] {res.get('message')}")


def cli_domain_browser():
    console.print(Panel(
        f"[bold cyan]🌐 Nova Astra 1000+ Domain Specializations Hub[/bold cyan]\n"
        f"[dim]Total Available: {domain_registry.total_count:,} domains across {len(domain_registry.get_all_categories())} categories.[/dim]",
        border_style="cyan"
    ))
    search_q = Prompt.ask("[bold yellow]Enter keyword to search domains (or press Enter to list categories)[/bold yellow]", default="").strip()
    if search_q:
        results = domain_registry.search_domains(search_q, limit=12)
        if not results:
            console.print("[yellow]No exact match found. Showing closest matches.[/yellow]")
        table = Table(title=f"Domain Search Results for '{search_q}'", border_style="magenta")
        table.add_column("ID", style="dim")
        table.add_column("Domain Specialization", style="bold cyan")
        table.add_column("Category", style="green")
        table.add_column("Scope Summary")
        for d in results:
            table.add_row(d.id, d.name, d.category, d.description[:60] + "...")
        console.print(table)
    else:
        cats = domain_registry.get_all_categories()
        table = Table(title="1,000+ Domain Categories", border_style="cyan")
        table.add_column("#", style="dim")
        table.add_column("Category", style="bold white")
        table.add_column("Domains Count", style="bold yellow")
        for i, c in enumerate(cats, 1):
            count = len(domain_registry.get_domains_by_category(c))
            table.add_row(str(i), c, f"{count} domains")
        console.print(table)
        console.print("[dim]Tip: Type '/domain <id_or_keyword>' to lock into any domain, or '/domain auto' to restore auto-detect.[/dim]")


def cli_main():
    reasoning_mode = False
    grounding_mode = False
    active_domain = None
    active_domain_name = None
    print_banner(reasoning_mode, grounding_mode, active_domain_name)

    with console.status("[bold green]🌟 Initializing Nova Astra Engine...[/bold green]"):
        try:
            engine.load_model()
        except Exception as e:
            console.print(f"[bold red]Initialization failed:[/bold red] {e}")
            return

    system_prompt = SYSTEM_PROMPT_PRESETS["ChatGPT Astra Core (Omniscient & Empathetic)"]
    messages: List[Dict[str, str]] = []

    while True:
        try:
            user_input = Prompt.ask(f"\n[bold green]👤 You[/bold green] [dim]({token_vault.get_balance()} ASTRA)[/dim]").strip()
        except (KeyboardInterrupt, EOFError):
            console.print("\n[bold yellow]Exiting Nova AI. Farewell![/bold yellow]")
            break

        if not user_input:
            continue

        if user_input.startswith("/"):
            cmd = user_input.lower().split()[0]
            if cmd in ("/quit", "/exit"):
                console.print("[bold yellow]Exiting Nova AI. Farewell![/bold yellow]")
                break
            elif cmd == "/help":
                print_help()
                continue
            elif cmd == "/tokens":
                show_tokens()
                continue
            elif cmd == "/domains":
                cli_domain_browser()
                continue
            elif cmd == "/domain":
                parts = user_input.split(maxsplit=1)
                if len(parts) < 2 or parts[1].strip().lower() in ("auto", "none", "reset"):
                    active_domain = None
                    active_domain_name = None
                    console.print("[bold green]✓ Domain specialization reset to: AUTO-DETECT (1,425 domains active).[/bold green]")
                else:
                    query = parts[1].strip()
                    matched = domain_registry.get_domain(query)
                    if not matched:
                        matches = domain_registry.search_domains(query, limit=1)
                        if matches:
                            matched = matches[0]
                    if matched:
                        active_domain = matched.id
                        active_domain_name = f"{matched.name} ({matched.category})"
                        console.print(f"[bold green]✓ Specialization locked to:[/bold green] [bold cyan]{active_domain_name}[/bold cyan]")
                    else:
                        console.print(f"[bold red]Domain '{query}' not found. Type '/domains' to search 1,000+ domains.[/bold red]")
                continue
            elif cmd == "/train":
                cli_fine_tuning()
                continue
            elif cmd in ("/search", "/websearch"):
                cli_web_search()
                continue
            elif cmd == "/ground":
                grounding_mode = not grounding_mode
                status_str = "ENABLED" if grounding_mode else "DISABLED"
                console.print(f"[bold cyan]🌐 Real-Time Web Grounding:[/bold cyan] {status_str}")
                continue
            elif cmd == "/image":
                cli_image_generator()
                continue
            elif cmd == "/web":
                cli_website_generator()
                continue
            elif cmd == "/think":
                reasoning_mode = not reasoning_mode
                status_str = "ENABLED (Chain-of-Thought active)" if reasoning_mode else "DISABLED"
                console.print(f"[bold cyan]🧠 Astra Deep Reasoning Mode:[/bold cyan] {status_str}")
                continue
            elif cmd == "/spark":
                raw_query = Prompt.ask("[bold cyan]Enter query to enhance[/bold cyan]")
                enhanced = enhance_prompt_with_astra(raw_query)
                console.print(Panel(enhanced, title="[bold magenta]✨ Astra Spark[/bold magenta]", border_style="magenta"))
                continue
            elif cmd == "/clear":
                messages.clear()
                console.print("[bold yellow]✓ Conversation cleared.[/bold yellow]")
                continue
            elif cmd == "/stats":
                hw = get_system_hardware_info()
                table = Table(title="Astra Runtime Telemetry", border_style="cyan")
                table.add_column("Metric", style="bold")
                table.add_column("Value")
                table.add_row("Active Model", str(engine.model_id))
                table.add_row("Astra Balance", f"{token_vault.get_balance()} ASTRA")
                table.add_row("Device", engine.device.upper())
                table.add_row("RAM", f"{hw.get('available_ram_gb')} GB / {hw.get('total_ram_gb')} GB")
                console.print(table)
                continue
            else:
                console.print(f"[bold red]Unknown command:[/bold red] {user_input}. Type [bold yellow]/help[/bold yellow].")
                continue

        # Check and deduct tokens for chat
        action = "deep_reasoning" if reasoning_mode else "chat_query"
        total_cost = TOKEN_COSTS[action] + (TOKEN_COSTS["web_search"] if grounding_mode else 0)

        if token_vault.get_balance() < total_cost:
            console.print(f"[bold red]Insufficient tokens. Requires {total_cost} ASTRA, but you have {token_vault.get_balance()} ASTRA.[/bold red]")
            continue

        success, msg, remaining = token_vault.deduct_tokens(action)
        if not success:
            console.print(f"[bold red]{msg}[/bold red]")
            continue

        grounded_context = ""
        citations_summary = []
        if grounding_mode:
            token_vault.deduct_tokens("web_search")
            with console.status("[bold cyan]Querying live web sources...[/bold cyan]"):
                search_res = search_live_web(user_input, max_results=3, deduct_tokens=False)
            if search_res.get("status") == "success" and search_res.get("results"):
                grounded_context = search_res.get("grounding_context", "")
                citations_summary = search_res.get("results", [])

        messages.append({"role": "user", "content": user_input})
        gen_messages = []
        if system_prompt:
            gen_messages.append({"role": "system", "content": system_prompt})

        for m in messages[:-1]:
            gen_messages.append({"role": m["role"], "content": m["content"]})

        if grounded_context:
            gen_messages.append({"role": "user", "content": f"{grounded_context}\n\nUser Question: {user_input}"})
        else:
            gen_messages.append({"role": "user", "content": user_input})

        gen_cfg = GenerationConfig(
            max_new_tokens=config.gen_config.max_new_tokens,
            temperature=config.gen_config.temperature,
            deep_reasoning=reasoning_mode,
            domain_id=active_domain,
            auto_detect_domain=(active_domain is None)
        )

        rem_bal = token_vault.get_balance()
        ground_tag = " | 🌐 Web Grounded" if grounding_mode else ""
        domain_tag = f" | 🎓 {active_domain_name}" if active_domain_name else ""
        console.print(f"\n[bold magenta]🌟 Nova Astra:[/bold magenta] [dim](-{total_cost} ASTRA{ground_tag}{domain_tag} | Bal: {rem_bal})[/dim]")
        full_response = ""
        last_metrics = {}

        try:
            for partial_text, metrics in engine.stream_chat(gen_messages, gen_cfg):
                new_chunk = partial_text[len(full_response):]
                console.print(new_chunk, end="", style="white")
                full_response = partial_text
                last_metrics = metrics

            console.print()

            if citations_summary:
                console.print("\n[bold cyan]🌐 Verified Live Sources:[/bold cyan]")
                for idx, r in enumerate(citations_summary, 1):
                    console.print(f"  [cyan][{idx}][/cyan] {r['title']} [dim]({r['url']})[/dim]")

            messages.append({"role": "assistant", "content": full_response})

            tps = last_metrics.get("tokens_per_second", 0.0)
            ttft = last_metrics.get("time_to_first_token_ms", 0.0)
            tokens = last_metrics.get("tokens_generated", 0)
            lat = last_metrics.get("total_latency_seconds", 0.0)
            detected_domain = last_metrics.get("active_domain")
            domain_info = f" | 🎓 Specialization: {detected_domain}" if detected_domain else ""
            console.print(
                f"[dim]⚡ Speed: {tps} tps | TTFT: {ttft}ms | Tokens: {tokens} | Latency: {lat}s{domain_info} | Tokens Rem: {rem_bal} ASTRA[/dim]"
            )

        except Exception as e:
            console.print(f"\n[bold red]Error during generation:[/bold red] {e}")


if __name__ == "__main__":
    cli_main()