import asyncio
import json
import time
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4

from zoneinfo import ZoneInfo

from claude_agent_sdk import query, ClaudeAgentOptions, AssistantMessage, TextBlock, ResultMessage

from mixersystem.data.repository import _ctx_var, parse_result, RETRY_FORMAT_PROMPT, get_project_root
from mixersystem.data.trace_db import insert_trace_event

TZ_LA = ZoneInfo("America/Los_Angeles")
TZ_YER = ZoneInfo("Asia/Yerevan")
PROVIDER_CALL_TIMEOUT_SECONDS = 15 * 60

# Model tiers: agents declare S / M / L, resolved to provider-specific model names.
_MODEL_MAP = {
    "claude": {"L": "opus", "M": "sonnet", "S": "haiku"},
    "gemini": {"L": "gemini-3-pro-preview", "M": "gemini-3-flash-preview", "S": "gemini-2.5-flash"},
    "codex":  {"L": "gpt-5.3-codex", "M": "gpt-5.3-codex", "S": "gpt-5.1-codex-mini"},
}

# Codex reasoning effort per tier.
_CODEX_REASONING_EFFORT = {"L": "xhigh", "M": "medium", "S": "high"}


def _resolve_model(tier: str, provider: str) -> str:
    """Resolve a model tier (S/M/L) to a provider-specific model name."""
    provider_map = _MODEL_MAP.get(provider)
    if not provider_map:
        raise ValueError(f"Unsupported provider: {provider}")
    model = provider_map.get(tier)
    if not model:
        raise ValueError(f"Unknown model tier '{tier}' for provider '{provider}'. Use S, M, or L.")
    return model


def _timestamp() -> str:
    now = datetime.now(tz=timezone.utc)
    utc = now.strftime("%H:%M:%S")
    la_time = now.astimezone(TZ_LA)
    yer_time = now.astimezone(TZ_YER)
    la = la_time.strftime(f"%H:%M:%S {la_time.tzname()}")
    yer = yer_time.strftime(f"%H:%M:%S {yer_time.tzname()}")
    return f"{la} / {yer} / {utc} UTC"


def _fmt_duration(seconds: float) -> str:
    if seconds < 60:
        return f"{seconds:.1f}s"
    m, s = divmod(int(seconds), 60)
    return f"{m}m {s}s"


def _fmt_cost(cost: float | None) -> str:
    if cost is None:
        return "n/a"
    return f"${cost:.4f}"


def _ensure_logs_dir(session_folder: str):
    logs_dir = Path(session_folder) / "logs"
    logs_dir.mkdir(parents=True, exist_ok=True)
    return logs_dir


async def _communicate_with_timeout(*, proc: asyncio.subprocess.Process,
                                    timeout_seconds: int,
                                    tool_name: str) -> tuple[bytes, bytes]:
    try:
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout_seconds)
    except asyncio.TimeoutError as exc:
        proc.kill()
        await proc.communicate()
        raise TimeoutError(
            f"{tool_name} timed out after {timeout_seconds}s."
        ) from exc

    if proc.returncode != 0:
        stderr_text = stderr.decode(errors="replace").strip()
        stdout_text = stdout.decode(errors="replace").strip()
        detail = stderr_text or stdout_text or "(no output)"
        raise RuntimeError(
            f"{tool_name} exited with code {proc.returncode}: {detail[:400]}"
        )

    return stdout, stderr


def _log_trace(session_folder: str, agent_name: str, depth: int, entry_type: str,
               call_id: str = None,
               model: str = None, session_id: str = None, is_resume: bool = False,
               duration: float = None, cost: float = None, status: str = None):
    """Light log — agent_trace.log (flat). No prompts, no responses."""
    logs_dir = _ensure_logs_dir(session_folder)
    trace_log_path = logs_dir / "agent_trace.log"

    timestamp = _timestamp()
    mode = "resume" if is_resume else "call"

    if entry_type == "start":
        call_tag = f" ({call_id[:8]})" if call_id else ""
        flat = f"[{timestamp}] START {agent_name}{call_tag} ({mode}) model={model or 'sonnet'}\n"
        with open(trace_log_path, "a") as f:
            f.write(flat)
        return

    if entry_type == "end":
        parts = [f"session={session_id or 'unknown'}"]
        if duration is not None:
            parts.append(f"duration={_fmt_duration(duration)}")
        if cost is not None:
            parts.append(f"cost={_fmt_cost(cost)}")
        if status:
            parts.append(f"status={status}")
        info = " | ".join(parts)
        call_tag = f" ({call_id[:8]})" if call_id else ""
        flat_end = f"[{timestamp}] END {agent_name}{call_tag} {info}\n"

    elif entry_type == "error":
        call_tag = f" ({call_id[:8]})" if call_id else ""
        flat_end = f"[{timestamp}] ERROR {agent_name}{call_tag} {status or ''}\n"

    else:
        return

    with open(trace_log_path, "a") as f:
        f.write(flat_end)


def _log_raw(session_folder: str, agent_name: str, depth: int, entry_type: str,
             model: str = None, prompt: str = None, response: str = None,
             session_id: str = None, is_resume: bool = False,
             duration: float = None, cost: float = None):
    """Full log — agent_raw.log. Includes prompts and responses."""
    logs_dir = _ensure_logs_dir(session_folder)
    raw_path = logs_dir / "agent_raw.log"

    timestamp = _timestamp()
    mode = "resume" if is_resume else "call"

    if entry_type == "start":
        entry = f"[{timestamp}] START {agent_name} ({mode}) model={model or 'sonnet'}\n{prompt or ''}\n\n"

    elif entry_type == "end":
        meta_parts = [f"session={session_id or 'unknown'}"]
        if duration is not None:
            meta_parts.append(f"duration={_fmt_duration(duration)}")
        if cost is not None:
            meta_parts.append(f"cost={_fmt_cost(cost)}")
        meta = " | ".join(meta_parts)
        entry = f"[{timestamp}] END {agent_name} {meta}\n{response or ''}\n\n"

    elif entry_type == "error":
        entry = f"[{timestamp}] ERROR {agent_name}\n{response or ''}\n\n"

    else:
        return

    with open(raw_path, "a") as f:
        f.write(entry)


def _log_db(session_folder: str, agent_name: str, depth: int, event_type: str,
            call_id: str = "", model: str = "", provider: str = "",
            session_id: str = "", is_resume: bool = False,
            duration: float = None, cost: float = None,
            status: str = "", error_message: str = ""):
    """Write a trace event to the per-session SQLite DB.  Never crashes the agent."""
    try:
        mode = "resume" if is_resume else "call"
        insert_trace_event(
            session_folder,
            call_id=call_id,
            agent_name=agent_name,
            event_type=event_type,
            model=model,
            provider=provider,
            mode=mode,
            session_id=session_id or "",
            depth=depth,
            duration_secs=duration,
            cost_usd=cost,
            status=status or "",
            error_message=error_message,
            is_resume=is_resume,
        )
    except Exception:
        pass


async def _call_claude(prompt: str, model: str, agent_name: str,
                       permission_mode: str, settings: str,
                       is_resume: bool, session_id: str,
                       depth: int = 0,
                       project_root: Path | None = None) -> tuple[str, str, float | None]:
    """Call Claude Agent SDK. Returns (response_text, session_id, cost_usd)."""
    opts = {"permission_mode": permission_mode, "model": model}

    if is_resume and session_id:
        opts = {"permission_mode": permission_mode, "resume": session_id}

    if not is_resume:
        opts["cwd"] = str(project_root or get_project_root())

    if settings:
        opts["settings"] = settings

    opts["disallowed_tools"] = ["AskUserQuestion"]

    claude_options = ClaudeAgentOptions(**opts)

    async def _run_query() -> tuple[str, str | None, float | None]:
        response_text = ""
        result_session_id = None
        cost_usd = None

        async for message in query(prompt=prompt, options=claude_options):
            if isinstance(message, AssistantMessage):
                for block in message.content:
                    if isinstance(block, TextBlock):
                        if response_text:
                            response_text += "\n"
                        response_text += block.text
            elif isinstance(message, ResultMessage):
                result_session_id = message.session_id
                cost_usd = message.total_cost_usd

        return response_text, result_session_id, cost_usd

    try:
        return await asyncio.wait_for(
            _run_query(),
            timeout=PROVIDER_CALL_TIMEOUT_SECONDS,
        )
    except asyncio.TimeoutError as exc:
        raise TimeoutError(
            f"Claude call timed out after {PROVIDER_CALL_TIMEOUT_SECONDS}s."
        ) from exc


async def _call_gemini(prompt: str, model: str, agent_name: str,
                       is_resume: bool, session_id: str,
                       approval_mode: str = "yolo",
                       project_root: Path | None = None) -> tuple[str, str, float | None]:
    """Call Gemini CLI. Returns (response_text, session_id, cost_usd).
    approval_mode: 'yolo' (auto-approve all), 'plan' (read-only), 'auto_edit', 'default'
    """
    cmd = ["gemini", "-m", model, "-p", prompt, "--approval-mode", approval_mode, "-o", "json"]

    if is_resume and session_id:
        cmd = ["gemini", "-m", model, "--resume", session_id, "-p", prompt, "--approval-mode", approval_mode, "-o", "json"]

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        cwd=str(project_root or get_project_root()),
    )
    stdout, stderr = await _communicate_with_timeout(
        proc=proc,
        timeout_seconds=PROVIDER_CALL_TIMEOUT_SECONDS,
        tool_name="Gemini CLI",
    )

    raw = stdout.decode(errors="replace")

    # Parse JSON response (banner goes to stderr, stdout should be clean JSON)
    json_start = raw.find("{")
    if json_start == -1:
        stderr_text = stderr.decode(errors="replace")
        raise RuntimeError(f"Gemini CLI returned no JSON. stdout: {raw[:200]}, stderr: {stderr_text[:200]}")

    data = json.loads(raw[json_start:])

    if "error" in data:
        err = data["error"]
        raise RuntimeError(f"Gemini CLI error: {err.get('message', err)}")

    response_text = data.get("response", "")
    result_session_id = data.get("session_id")

    return response_text, result_session_id, None


async def _call_codex(prompt: str, model: str, agent_name: str,
                      is_resume: bool, session_id: str,
                      reasoning_effort: str = "medium",
                      project_root: Path | None = None) -> tuple[str, str, float | None]:
    """Call Codex CLI. Returns (response_text, session_id, cost_usd)."""
    if is_resume and session_id:
        cmd = ["codex", "exec", "--full-auto", "-c", f"model_reasoning_effort={reasoning_effort}", "-m", model, "resume", session_id, prompt, "--json"]
    else:
        cmd = ["codex", "exec", "--full-auto", "-c", f"model_reasoning_effort={reasoning_effort}", "-m", model, prompt, "--json"]

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        cwd=str(project_root or get_project_root()),
    )
    stdout, stderr = await _communicate_with_timeout(
        proc=proc,
        timeout_seconds=PROVIDER_CALL_TIMEOUT_SECONDS,
        tool_name="Codex CLI",
    )

    raw = stdout.decode(errors="replace")

    # Parse JSONL — extract thread_id and last agent_message
    response_text = ""
    result_session_id = None

    for line in raw.strip().splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            event = json.loads(line)
        except json.JSONDecodeError:
            continue

        if event.get("type") == "thread.started":
            result_session_id = event.get("thread_id")
        elif event.get("type") == "item.completed":
            item = event.get("item", {})
            if item.get("type") == "agent_message":
                response_text = item.get("text", "")

    if not response_text:
        stderr_text = stderr.decode(errors="replace")
        raise RuntimeError(
            f"Codex CLI returned no agent_message. stdout: {raw[:200]}, stderr: {stderr_text[:200]}"
        )

    return response_text, result_session_id, None


async def _dispatch(prompt, model, provider, agent_name,
                    permission_mode, settings, approval_mode,
                    is_resume, session_id, depth=0,
                    project_root: Path | None = None,
                    reasoning_effort: str = "medium") -> tuple[str, str, float | None]:
    """Dispatch to the right provider. Returns (response_text, session_id, cost_usd)."""
    if provider == "gemini":
        return await _call_gemini(
            prompt, model, agent_name, is_resume, session_id, approval_mode,
            project_root=project_root,
        )
    if provider == "codex":
        return await _call_codex(
            prompt, model, agent_name, is_resume, session_id,
            reasoning_effort=reasoning_effort,
            project_root=project_root,
        )
    if provider == "claude":
        return await _call_claude(
            prompt, model, agent_name,
            permission_mode, settings, is_resume, session_id, depth,
            project_root=project_root,
        )
    raise ValueError(
        f"Unsupported provider: {provider}. "
        "Supported providers: claude, gemini, codex."
    )


def _context_label(session_folder: str) -> str:
    """Derive a short context label from the session folder path (e.g. 'plan_a' for a branch)."""
    if not session_folder:
        return ""
    name = Path(session_folder).name
    parent = Path(session_folder).parent.name
    if parent == "branches":
        return f" {name}"
    return ""


def _print_status(agent_name: str, context: str, message: str) -> None:
    """Print a status line to stderr for terminal observability."""
    import sys
    print(f"[{agent_name}]{context} {message}", file=sys.stderr, flush=True)


async def call_agent(prompt: str, model: str = "M", provider: str = "claude",
                     agent_name: str = "agent",
                     permission_mode: str = "bypassPermissions",
                     settings: str = None, approval_mode: str = "yolo",
                     is_resume: bool = False, session_id: str = None) -> tuple[str, str]:
    """
    Call an LLM and return (response_text, session_id).
    model: "S", "M", or "L" tier — resolved to provider-specific model name.
    provider: "claude", "gemini", or "codex"
    approval_mode: Gemini only — 'yolo', 'plan', 'auto_edit', 'default'
    Always retries once when response is missing [RESULT] block.
    """
    ctx = _ctx_var.get()
    if ctx.lite:
        model = "S"
    session_folder = ctx.session_folder
    depth = ctx.depth
    project_root = Path(ctx.project_root).resolve() if ctx.project_root else get_project_root()
    resolved_model = _resolve_model(model, provider)
    reasoning_effort = _CODEX_REASONING_EFFORT.get(model, "medium") if provider == "codex" else "medium"
    display_model = f"{provider}/{resolved_model}" if provider != "claude" else resolved_model
    call_id = uuid4().hex
    ctx_label = _context_label(session_folder)
    resume_tag = " resume" if is_resume else ""

    # Log start
    if session_folder:
        _log_trace(session_folder, agent_name, depth, "start",
                   call_id=call_id, model=display_model, is_resume=is_resume)
        _log_raw(session_folder, agent_name, depth, "start",
                 model=display_model, prompt=prompt, is_resume=is_resume)
        _log_db(session_folder, agent_name, depth, "start",
                call_id=call_id, model=display_model, provider=provider,
                is_resume=is_resume)

    _print_status(agent_name, ctx_label, f"starting ({display_model}){resume_tag} ...")

    t0 = time.monotonic()
    cost_usd = None

    try:
        response_text, result_session_id, cost_usd = await _dispatch(
            prompt, resolved_model, provider, agent_name,
            permission_mode, settings, approval_mode,
            is_resume, session_id, depth,
            project_root=project_root,
            reasoning_effort=reasoning_effort,
        )

        # Retry once if [RESULT] block is missing
        if "[RESULT]" not in response_text:
            duration = time.monotonic() - t0
            _print_status(agent_name, ctx_label, f"retrying — missing [RESULT] block ({_fmt_duration(duration)})")
            if session_folder:
                _log_trace(session_folder, agent_name, depth, "end",
                           call_id=call_id, session_id=result_session_id, duration=duration, cost=cost_usd)
                _log_raw(session_folder, agent_name, depth, "end",
                         response=response_text, session_id=result_session_id, duration=duration, cost=cost_usd)
                _log_db(session_folder, agent_name, depth, "end",
                        call_id=call_id, provider=provider,
                        session_id=result_session_id or "", duration=duration, cost=cost_usd)
                _log_trace(session_folder, agent_name, depth, "start",
                           call_id=call_id, model=display_model, is_resume=True)
                _log_raw(session_folder, agent_name, depth, "start",
                         model=display_model, prompt=RETRY_FORMAT_PROMPT, is_resume=True)
                _log_db(session_folder, agent_name, depth, "start",
                        call_id=call_id, model=display_model, provider=provider,
                        is_resume=True)
            t0 = time.monotonic()
            response_text, result_session_id, cost_usd = await _dispatch(
                RETRY_FORMAT_PROMPT, resolved_model, provider, agent_name,
                permission_mode, settings, approval_mode,
                True, result_session_id, depth,
                project_root=project_root,
                reasoning_effort=reasoning_effort,
            )

    except Exception as e:
        error_msg = f"{type(e).__name__}: {e}"
        duration = time.monotonic() - t0
        _print_status(agent_name, ctx_label, f"error ({_fmt_duration(duration)})")
        if session_folder:
            _log_trace(session_folder, agent_name, depth, "error",
                       call_id=call_id, status=error_msg, duration=duration)
            _log_raw(session_folder, agent_name, depth, "error",
                     response=error_msg, duration=duration)
            _log_db(session_folder, agent_name, depth, "error",
                    call_id=call_id, provider=provider,
                    error_message=error_msg, duration=duration)
        raise

    duration = time.monotonic() - t0

    # Parse status from response for trace summary
    status = None
    if "[RESULT]" in response_text:
        result = parse_result(response_text)
        status = result.status or None

    status_tag = f" — {status}" if status else ""
    _print_status(agent_name, ctx_label, f"done ({_fmt_duration(duration)}){status_tag}")

    # Log end
    if session_folder:
        _log_trace(session_folder, agent_name, depth, "end",
                   call_id=call_id, session_id=result_session_id, duration=duration,
                   cost=cost_usd, status=status, is_resume=is_resume)
        _log_raw(session_folder, agent_name, depth, "end",
                 response=response_text, session_id=result_session_id,
                 duration=duration, cost=cost_usd, is_resume=is_resume)
        _log_db(session_folder, agent_name, depth, "end",
                call_id=call_id, provider=provider,
                session_id=result_session_id or "", duration=duration,
                cost=cost_usd, status=status or "", is_resume=is_resume)

    return response_text, result_session_id
