"""AI client: context assembly + streaming response to WebSocket viewers.

Supports Claude (Anthropic) and Gemini (Google) as providers.
Provider is auto-detected from available API keys (Claude preferred).
"""

from __future__ import annotations

import base64
import logging
from pathlib import Path
from typing import TYPE_CHECKING, List

if TYPE_CHECKING:
    from ai_interview.config import Config
    from ai_interview.state import AppState

logger = logging.getLogger(__name__)


def _assemble_messages(state: "AppState", config: "Config") -> list:
    """Build the messages list for the API call, including conversation history."""
    transcript = state.transcript_buffer.get_recent(minutes=config.context_minutes) if state.transcript_buffer is not None else ""
    # Clear transcript after consuming — conversation_history carries prior context,
    # so the next query should only see what's spoken AFTER this point
    if state.transcript_buffer is not None:
        state.transcript_buffer.clear()
    screenshots: List[Path] = state.screenshot_buffer.pop_all() if state.screenshot_buffer is not None else []

    # Read clipboard only when explicitly requested (Ctrl double-tap sets the flag)
    clipboard = ""
    if state.include_clipboard:
        state.include_clipboard = False
        try:
            import subprocess
            result = subprocess.run(["pbpaste"], capture_output=True, text=True, timeout=1)
            clipboard = result.stdout.strip()
        except Exception:
            pass

    content = []

    if transcript:
        content.append({
            "type": "text",
            "text": f"## Interview context (last {config.context_minutes} minutes of audio):\n\n{transcript}",
        })

    if clipboard:
        content.append({
            "type": "text",
            "text": f"## Additional context (from clipboard):\n\n{clipboard}",
        })
        logger.debug("Clipboard included (%d chars)", len(clipboard))

    if screenshots:
        from ai_interview.screenshot import encode_screenshot_b64
        for path in screenshots:
            try:
                b64 = encode_screenshot_b64(path)
                content.append({
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/jpeg",
                        "data": b64,
                    },
                })
                logger.debug("Attached screenshot: %s", path.name)
            except Exception as exc:
                logger.warning("Could not attach screenshot %s: %s", path.name, exc)

    if not content:
        content.append({
            "type": "text",
            "text": (
                "No audio transcript or screenshots captured yet. "
                "Please describe the coding challenge you are working on."
            ),
        })

    # Prepend conversation history, then add current turn
    messages = list(state.conversation_history)
    messages.append({"role": "user", "content": content})
    return messages


# ---------------------------------------------------------------------------
# Claude (Anthropic) streaming
# ---------------------------------------------------------------------------

async def _stream_claude(state: "AppState", config: "Config", messages: list, current_user_turn: dict) -> None:
    """Stream response from Claude API."""
    from ai_interview.server.websocket import broadcast
    import anthropic

    is_oauth = config.anthropic_api_key.startswith("sk-ant-oat")
    if is_oauth:
        client = anthropic.AsyncAnthropic(
            api_key="placeholder",
            default_headers={"Authorization": f"Bearer {config.anthropic_api_key}"},
        )
    else:
        client = anthropic.AsyncAnthropic(api_key=config.anthropic_api_key)

    full_text = ""

    try:
        async with client.messages.stream(
            model=config.claude_model,
            max_tokens=config.max_tokens,
            system=config.system_prompt,
            messages=messages,
        ) as stream:
            async for text_delta in stream.text_stream:
                full_text += text_delta
                await broadcast({"type": "chunk", "text": text_delta})

        state.last_response = full_text
        state.conversation_history.append(current_user_turn)
        state.conversation_history.append({"role": "assistant", "content": full_text})

        max_entries = config.max_history_turns * 2
        if len(state.conversation_history) > max_entries:
            state.conversation_history = state.conversation_history[-max_entries:]

        await broadcast({"type": "done", "full_text": full_text})
        logger.info("Claude query complete (%d chars)", len(full_text))

    except anthropic.AuthenticationError:
        msg = "Invalid Anthropic API key. Check your --api-key value."
        logger.error(msg)
        await broadcast({"type": "error", "message": msg})

    except anthropic.BadRequestError as exc:
        if "token" in str(exc).lower():
            trimmed = len(state.conversation_history)
            half = max(2, len(state.conversation_history) // 2)
            state.conversation_history = state.conversation_history[half:]
            logger.warning("Context too long (%d entries) — trimmed to %d, retrying", trimmed, len(state.conversation_history))
            await broadcast({"type": "start"})
            await broadcast({"type": "chunk", "text": "_Context too long — trimming history and retrying…_"})

            messages_retry = list(state.conversation_history)
            messages_retry.append(current_user_turn)
            try:
                full_text = ""
                async with client.messages.stream(
                    model=config.claude_model,
                    max_tokens=config.max_tokens,
                    system=config.system_prompt,
                    messages=messages_retry,
                ) as stream:
                    await broadcast({"type": "start"})
                    async for text_delta in stream.text_stream:
                        full_text += text_delta
                        await broadcast({"type": "chunk", "text": text_delta})

                state.last_response = full_text
                state.conversation_history.append(current_user_turn)
                state.conversation_history.append({"role": "assistant", "content": full_text})
                await broadcast({"type": "done", "full_text": full_text})
            except Exception as retry_exc:
                logger.error("Retry after trim failed: %s", retry_exc)
                await broadcast({"type": "error", "message": f"Failed after trimming: {retry_exc}"})
        else:
            logger.error("Bad request: %s", exc)
            await broadcast({"type": "error", "message": f"Bad request: {exc}"})

    except anthropic.APIConnectionError as exc:
        logger.error("Connection error: %s", exc)
        await broadcast({"type": "error", "message": f"Connection error: {exc}"})

    except Exception as exc:
        logger.error("Unexpected error: %s", exc)
        await broadcast({"type": "error", "message": f"Unexpected error: {exc}"})


# ---------------------------------------------------------------------------
# Gemini (Google) streaming
# ---------------------------------------------------------------------------

def _messages_to_gemini_contents(messages: list, system_prompt: str) -> list:
    """Convert internal message format to Gemini contents list."""
    from google.genai import types

    contents = []

    for msg in messages:
        role = "user" if msg["role"] == "user" else "model"
        parts = []

        if isinstance(msg["content"], str):
            parts.append(types.Part.from_text(text=msg["content"]))
        elif isinstance(msg["content"], list):
            for block in msg["content"]:
                if block.get("type") == "text":
                    parts.append(types.Part.from_text(text=block["text"]))
                elif block.get("type") == "image":
                    b64_data = block["source"]["data"]
                    mime = block["source"].get("media_type", "image/jpeg")
                    image_bytes = base64.b64decode(b64_data)
                    parts.append(types.Part.from_bytes(data=image_bytes, mime_type=mime))

        if parts:
            contents.append(types.Content(role=role, parts=parts))

    return contents


async def _stream_gemini(state: "AppState", config: "Config", messages: list, current_user_turn: dict) -> None:
    """Stream response from Gemini API."""
    import asyncio
    from ai_interview.server.websocket import broadcast

    try:
        from google import genai
        from google.genai import types

        client = genai.Client(api_key=config.google_api_key)
        contents = _messages_to_gemini_contents(messages, config.system_prompt)

        # Run synchronous streaming in a thread to avoid blocking asyncio
        full_text = ""

        def _stream_sync():
            nonlocal full_text
            chunks = []
            for chunk in client.models.generate_content_stream(
                model=config.gemini_model,
                contents=contents,
                config=types.GenerateContentConfig(
                    system_instruction=config.system_prompt,
                    max_output_tokens=config.max_tokens,
                ),
            ):
                if chunk.text:
                    full_text += chunk.text
                    chunks.append(chunk.text)
            return chunks

        loop = asyncio.get_event_loop()
        chunks = await loop.run_in_executor(None, _stream_sync)

        # Broadcast all chunks
        for chunk_text in chunks:
            await broadcast({"type": "chunk", "text": chunk_text})

        state.last_response = full_text
        state.conversation_history.append(current_user_turn)
        state.conversation_history.append({"role": "assistant", "content": full_text})

        max_entries = config.max_history_turns * 2
        if len(state.conversation_history) > max_entries:
            state.conversation_history = state.conversation_history[-max_entries:]

        await broadcast({"type": "done", "full_text": full_text})
        logger.info("Gemini query complete (%d chars)", len(full_text))

    except ImportError:
        msg = "google-genai not installed. Run: pip install google-genai"
        logger.error(msg)
        await broadcast({"type": "error", "message": msg})

    except Exception as exc:
        error_str = str(exc).lower()
        if "api key" in error_str or "authenticate" in error_str or "403" in error_str:
            msg = "Invalid Google API key. Check your key at https://aistudio.google.com/api-keys"
            logger.error(msg)
            await broadcast({"type": "error", "message": msg})
        elif "quota" in error_str or "429" in error_str or "rate" in error_str:
            msg = "Gemini rate limit reached (free tier: 10 req/min). Wait a moment and retry."
            logger.warning(msg)
            await broadcast({"type": "error", "message": msg})
        else:
            msg = f"Gemini error: {exc}"
            logger.error(msg)
            await broadcast({"type": "error", "message": msg})


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

async def handle_query(state: "AppState", config: "Config") -> None:
    """Assemble context and stream AI response to all connected viewers.

    Auto-selects provider based on available API keys (Claude preferred).
    """
    from ai_interview.server.websocket import broadcast

    if state.status == "processing":
        logger.debug("Query already in progress, ignoring duplicate trigger")
        return

    state.status = "processing"
    await broadcast({"type": "status", "value": "processing"})
    await broadcast({"type": "start"})

    messages = _assemble_messages(state, config)
    current_user_turn = messages[-1]

    try:
        provider = config.provider
        if provider == "claude":
            await _stream_claude(state, config, messages, current_user_turn)
        elif provider == "gemini":
            await _stream_gemini(state, config, messages, current_user_turn)
        else:
            await broadcast({"type": "error", "message": "No API key configured. Run: ai-interview setup"})
    finally:
        state.status = "listening" if state.audio_active else "idle"
        await broadcast({"type": "status", "value": state.status})
