"""
chat.py — Streaming chat with the OneFirewall AI Gateway.

Uses SSE (Server-Sent Events) streaming from /api/v1/chat/completions.
Session history is persisted to ~/.onefirewall-ai/sessions/session1.md.
"""

import json
import sys
from datetime import datetime
from typing import Generator

import requests
from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel

from onefirewall_ai.config import (
    load_access,
    is_logged_in,
    SOUL_FILE,
    MEMORY_FILE,
    get_session_file,
    create_session_file,
)

console = Console()

# ─── Session history ──────────────────────────────────────────────────────────

def _load_session_messages(session_id: str) -> list[dict]:
    """
    Parse a session file into a list of OpenAI-style messages.
    We look for **User:** / **Assistant:** markers written by _append_to_session.
    """
    messages: list[dict] = []
    path = get_session_file(session_id)
    if not path.exists():
        return messages

    text = path.read_text(encoding="utf-8")
    for block in text.split("---"):
        lines = block.strip().splitlines()
        user_msg = ""
        asst_msg = ""
        for line in lines:
            if line.startswith("**User:** "):
                user_msg = line[len("**User:** "):]
            elif line.startswith("**Assistant:** "):
                asst_msg = line[len("**Assistant:** "):]
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if asst_msg:
            messages.append({"role": "assistant", "content": asst_msg})
    return messages

def _append_to_session(session_id: str, user_msg: str, assistant_msg: str) -> None:
    """Append a conversation turn to the session file."""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
    entry = (
        f"\n## [{timestamp}]\n"
        f"**User:** {user_msg}\n"
        f"**Assistant:** {assistant_msg}\n\n"
        f"---\n"
    )
    path = get_session_file(session_id)
    with open(path, "a", encoding="utf-8") as f:
        f.write(entry)

# ─── Streaming ────────────────────────────────────────────────────────────────

def _stream_chat(
    endpoint: str,
    api_key: str,
    model: str,
    messages: list[dict],
    web_search: bool = True,
) -> Generator[str, None, None]:
    """
    POST to /api/v1/chat/completions with stream=True.
    Yields content tokens as they arrive.

    Supports two response formats automatically:
      - text/plain   → raw chunked text (what onefirewall.ai returns)
      - text/event-stream → standard SSE  "data: {...}"  JSON packets
    """
    url = endpoint.rstrip("/") + "/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream, text/plain",
    }
    payload = {
        "model": model,
        "messages": messages,
        "stream": True,
        "web_search": web_search,
        "pii": "disabled",
    }

    try:
        with requests.post(
            url, headers=headers, json=payload, stream=True, timeout=60
        ) as resp:
            if resp.status_code != 200:
                yield f"\n[ERROR] HTTP {resp.status_code}: {resp.text[:300]}"
                return

            content_type = resp.headers.get("Content-Type", "")
            is_sse = "event-stream" in content_type

            if is_sse:
                for line in resp.iter_lines():
                    if not line:
                        continue
                    line_str = line.decode("utf-8") if isinstance(line, bytes) else line
                    line_str = line_str.strip()
                    if not line_str.startswith("data:"):
                        continue
                    payload_str = line_str[len("data:"):].strip()
                    if payload_str == "[DONE]":
                        return
                    try:
                        data = json.loads(payload_str)
                        token = (
                            data.get("choices", [{}])[0]
                            .get("delta", {})
                            .get("content", "")
                        )
                        if token:
                            yield token
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue
            else:
                for raw_chunk in resp.iter_content(chunk_size=None):
                    if not raw_chunk:
                        continue
                    chunk_str = raw_chunk.decode("utf-8") if isinstance(raw_chunk, bytes) else raw_chunk
                    if chunk_str:
                        yield chunk_str

    except requests.exceptions.ConnectionError:
        yield f"\n[ERROR] Cannot connect to '{endpoint}'. Run `onefirewall-ai login` to check your config."
    except requests.exceptions.Timeout:
        yield "\n[ERROR] Request timed out after 60 seconds."
    except Exception as exc:
        yield f"\n[ERROR] {exc}"



# ─── Public API ───────────────────────────────────────────────────────────────

def send(prompt: str, session_id: str = None) -> None:
    """
    Main entry: load config, build messages, stream response,
    render to terminal, and save to session file.
    """
    import threading
    import time
    import itertools

    if not is_logged_in():
        console.print(
            Panel(
                "[bold yellow]Not logged in.[/bold yellow]\n"
                "Run [bold cyan]onefirewall-ai login[/bold cyan] first.",
                border_style="yellow",
            )
        )
        sys.exit(1)

    acc = load_access()
    endpoint = acc["endpoint"]
    api_key = acc["api_key"]
    model = acc.get("model", "openai/gpt-4o")

    # Build system prompt from SOUL.md
    soul = ""
    if SOUL_FILE.exists():
        soul_raw = SOUL_FILE.read_text(encoding="utf-8").strip()
        soul_lines = [l for l in soul_raw.splitlines() if not l.startswith("#")]
        soul = "\n".join(soul_lines).strip()

    # Build full messages list
    messages: list[dict] = []
    
    web_search = True
    if soul:
        messages.append({"role": "system", "content": soul})
        # web_search = False
    
    if not session_id:
        session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    create_session_file(session_id)

    messages.extend(_load_session_messages(session_id))
    messages.append({"role": "user", "content": prompt})

    # ── Thread-based thinking spinner ───────────────────────────────────────
    _stop_spinner = threading.Event()

    def _spinner_thread() -> None:
        frames = ["⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"]
        # Fallback ASCII for terminals that can't handle unicode
        ascii_frames = ["|", "/", "-", "\\"]
        try:
            for frame in itertools.cycle(frames):
                if _stop_spinner.is_set():
                    break
                sys.stdout.write(f"\r  {frame} OneFirewall is thinking...")
                sys.stdout.flush()
                time.sleep(0.08)
        except Exception:
            for frame in itertools.cycle(ascii_frames):
                if _stop_spinner.is_set():
                    break
                sys.stdout.write(f"\r  {frame} OneFirewall is thinking...")
                sys.stdout.flush()
                time.sleep(0.1)

    spinner = threading.Thread(target=_spinner_thread, daemon=True)
    spinner.start()

    # ── Start streaming (this blocks until first token or error) ────────────
    gen = _stream_chat(endpoint, api_key, model, messages, web_search=web_search)
    first = next(gen, None)

    # ── Stop spinner, clear the spinner line ────────────────────────────────
    _stop_spinner.set()
    spinner.join(timeout=0.5)
    sys.stdout.write("\r" + " " * 50 + "\r")  # clear spinner line
    sys.stdout.flush()

    # ── Stream response to terminal ─────────────────────────────────────────

    collected_tokens: list[str] = []

    if first is not None:
        sys.stdout.write(first)
        sys.stdout.flush()
        collected_tokens.append(first)

    for token in gen:
        sys.stdout.write(token)
        sys.stdout.flush()
        collected_tokens.append(token)

    sys.stdout.write("\n\n")
    sys.stdout.flush()

    # ── Save to session ─────────────────────────────────────────────────────
    full_response = "".join(collected_tokens).strip()
    if full_response and not full_response.startswith("[ERROR]"):
        _append_to_session(session_id, prompt, full_response)

