"""
EPI CLI Chat - Interactive evidence querying with AI.

Allows users to ask natural language questions about their .epi evidence files.
"""

import json
import os
import warnings
from pathlib import Path
from typing import Optional

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


from epi_core.container import EPIContainer


console = Console()


def load_steps_from_epi(epi_path: Path) -> list:
    """Load steps from an .epi file."""
    import tempfile
    
    temp_dir = Path(tempfile.mkdtemp())
    extracted = EPIContainer.unpack(epi_path, temp_dir)
    
    steps_file = extracted / "steps.jsonl"
    if not steps_file.exists():
        return []
    
    steps = []
    with open(steps_file, 'r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                steps.append(json.loads(line))
    
    return steps


def chat(
    epi_file: Path = typer.Argument(..., help="Path to .epi file to chat with"),
    query: str = typer.Option(None, "--query", "-q", help="Single question (non-interactive mode)"),
    model: str = typer.Option("gemini-2.0-flash", "--model", "-m", help="Gemini model to use")
):
    """
    Chat with your evidence file using AI.
    
    Ask natural language questions about what happened in your recording.
    
    Examples:
        epi chat my_recording.epi                    # Interactive mode
        epi chat my_recording.epi -q "What happened?"  # Single question
    """
    # Resolve path
    if not epi_file.exists():
        # Try epi-recordings directory
        recordings_dir = Path("./epi-recordings")
        potential_path = recordings_dir / f"{epi_file.stem}.epi"
        if potential_path.exists():
            epi_file = potential_path
        else:
            console.print(f"[red]Error:[/red] File not found: {epi_file}")
            raise typer.Exit(1)
    
    # Detect which provider to use based on available API keys
    google_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
    openai_key = os.environ.get("OPENAI_API_KEY")

    if not google_key and not openai_key:
        console.print(Panel(
            "[yellow]No API key found![/yellow]\n\n"
            "Set one of:\n"
            "  [cyan]set OPENAI_API_KEY=sk-...[/cyan]         (Windows)\n"
            "  [cyan]export OPENAI_API_KEY=sk-...[/cyan]       (Mac/Linux)\n\n"
            "  [cyan]set GOOGLE_API_KEY=your-key-here[/cyan]   (Windows)\n"
            "  [cyan]export GOOGLE_API_KEY=your-key-here[/cyan] (Mac/Linux)\n\n"
            "Get a free Google key at: https://aistudio.google.com/app/apikey",
            title="[!] API Key Required",
            border_style="yellow"
        ))
        raise typer.Exit(1)

    # Load the .epi file
    console.print(f"\n[dim]Loading evidence from:[/dim] {epi_file}")

    try:
        manifest = EPIContainer.read_manifest(epi_file)
        steps = load_steps_from_epi(epi_file)
    except Exception as e:
        console.print(f"[red]Error loading .epi file:[/red] {e}")
        raise typer.Exit(1)

    # Build context (shared by both providers)
    context = f"""You are an expert assistant analyzing an EPI evidence recording file.

The recording contains cryptographically signed, tamper-proof evidence of an AI workflow execution.

Recording metadata:
- Created: {manifest.created_at}
- Goal: {manifest.goal or 'Not specified'}
- Command: {manifest.cli_command or 'Not specified'}
- Workflow ID: {manifest.workflow_id}
- Total steps: {len(steps)}

Here are the recorded steps (this is the timeline of events):
{json.dumps(steps[:50], indent=2, default=str)[:8000]}

When answering questions:
1. Be specific and cite step indices when relevant
2. Distinguish between LLM requests, responses, and other events
3. If asked about security, note that API keys are automatically redacted
4. Keep answers concise but informative
"""

    # ---- Provider: Gemini (preferred if key available) ----
    if google_key:
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                import google.generativeai as genai
                import google.api_core.exceptions as _goog_exc

            genai.configure(api_key=google_key)
            ai_model = genai.GenerativeModel(model)
            active_provider = "gemini"
        except ImportError:
            if not openai_key:
                console.print(Panel(
                    "[red]google-generativeai not installed and no OPENAI_API_KEY set.[/red]\n\n"
                    "Install Gemini:  [cyan]pip install google-generativeai[/cyan]\n"
                    "Or set:         [cyan]OPENAI_API_KEY=sk-...[/cyan]",
                    title="[X] Missing Dependency",
                    border_style="red"
                ))
                raise typer.Exit(1)
            google_key = None  # Fall through to OpenAI
        except Exception as e:
            console.print(f"[red]Error initialising Gemini:[/red] {e}")
            if not openai_key:
                raise typer.Exit(1)
            google_key = None  # Fall through to OpenAI

    # ---- Provider: OpenAI (fallback or primary if no Google key) ----
    if not google_key:
        try:
            from openai import OpenAI as _OpenAIClient
            _openai_client = _OpenAIClient(api_key=openai_key)
            _openai_model = "gpt-4o-mini"
            active_provider = "openai"
            _goog_exc = None  # not available

            class _OpenAIChatAdapter:
                """Minimal adapter so the chat loop below is provider-agnostic."""
                def __init__(self, client, model):
                    self._client = client
                    self._model = model
                    self._history = [{"role": "system", "content": context}]

                def send_message(self, prompt: str):
                    self._history.append({"role": "user", "content": prompt})
                    resp = self._client.chat.completions.create(
                        model=self._model,
                        messages=self._history,
                    )
                    text = resp.choices[0].message.content or ""
                    self._history.append({"role": "assistant", "content": text})

                    class _Resp:
                        pass
                    r = _Resp()
                    r.text = text
                    return r

            ai_model = _OpenAIChatAdapter(_openai_client, _openai_model)
            # Wrap send_message so we can call start_chat() the same way
            chat_session = ai_model
        except ImportError:
            console.print(Panel(
                "[red]openai package not installed.[/red]\n\n"
                "Install with: [cyan]pip install openai[/cyan]\n"
                "Or set:       [cyan]GOOGLE_API_KEY=...[/cyan] to use Gemini instead.",
                title="[X] Missing Dependency",
                border_style="red"
            ))
            raise typer.Exit(1)
        except Exception as e:
            console.print(f"[red]Error initialising OpenAI:[/red] {e}")
            raise typer.Exit(1)
    else:
        chat_session = ai_model.start_chat(history=[])

    # Display header
    console.print()
    console.print(Panel(
        f"[bold cyan]EPI Evidence Chat[/bold cyan]\n\n"
        f"[dim]File:[/dim]     {epi_file.name}\n"
        f"[dim]Steps:[/dim]    {len(steps)}\n"
        f"[dim]Provider:[/dim] {active_provider}  [dim]({model if active_provider == 'gemini' else _openai_model})[/dim]\n\n"
        f"Ask questions about this evidence recording.\n"
        f"Type [yellow]exit[/yellow] or [yellow]quit[/yellow] to end the session.",
        border_style="cyan"
    ))
    console.print()

    def _send(prompt: str) -> str:
        """Send a message and return the text response."""
        if active_provider == "gemini":
            full_prompt = f"{context}\n\nUser question: {prompt}"
            response = chat_session.send_message(full_prompt)
        else:
            response = chat_session.send_message(prompt)
        return response.text

    # Non-interactive mode: answer single question and exit
    if query:
        try:
            text = _send(query)
            console.print("[bold green]AI:[/bold green]")
            console.print(Markdown(text))
            return
        except Exception as e:
            console.print(f"[red]Error:[/red] {e}")
            raise typer.Exit(1)

    # Interactive chat loop
    while True:
        try:
            question = Prompt.ask("[bold cyan]You[/bold cyan]")
        except (KeyboardInterrupt, EOFError):
            console.print("\n[dim]Goodbye![/dim]")
            break

        if question.lower() in ('exit', 'quit', 'q'):
            console.print("[dim]Goodbye![/dim]")
            break

        if not question.strip():
            continue

        try:
            text = _send(question)
            console.print()
            console.print("[bold green]AI:[/bold green]")
            console.print(Markdown(text))
            console.print()
        except Exception as e:
            # Handle Gemini-specific quota errors gracefully
            if _goog_exc and isinstance(e, _goog_exc.ResourceExhausted):
                console.print(Panel(
                    "[yellow]API Quota Exceeded[/yellow]\n\n"
                    "You have hit the rate limit. Please wait a minute before trying again.",
                    title="[!] Rate Limit",
                    border_style="yellow"
                ))
            elif _goog_exc and isinstance(e, _goog_exc.NotFound):
                console.print(f"[red]Error:[/red] Model not found. Try a different model with --model.")
            else:
                console.print(f"[red]Error:[/red] {e}")
                console.print("[dim]Try asking a different question.[/dim]")
            console.print()



 