"""FastAPI server backing the local web GUI.

Wraps a single global :class:`LocalCodingAgent`, exposes JSON endpoints for
chat, slash commands, and saved sessions, and serves the static SPA.
"""

from __future__ import annotations

import asyncio
import json
import re
import time
from pathlib import Path
from threading import Lock
from typing import Any
from urllib import error, request

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from src.agent_runtime import LocalCodingAgent
from src.agent_slash_commands import get_slash_command_specs
from src.agent_types import (
    AgentPermissions,
    AgentRuntimeConfig,
    ModelConfig,
)
from src.bundled_skills import get_bundled_skills
from src.session_store import (
    DEFAULT_AGENT_SESSION_DIR,
    StoredAgentSession,
    load_agent_session,
)


STATIC_DIR = Path(__file__).resolve().parents[2] / 'frontend' / 'legacy-static'


# ---------------------------------------------------------------------------
# Agent state holder
# ---------------------------------------------------------------------------

class AgentState:
    """Holds the live agent instance plus a lock for serialized access."""

    def __init__(
        self,
        *,
        cwd: Path,
        model: str,
        base_url: str,
        api_key: str,
        allow_shell: bool,
        allow_write: bool,
        session_directory: Path,
    ) -> None:
        self.cwd = cwd.resolve()
        self.session_directory = session_directory
        self._lock = Lock()
        self._agents: dict[str, LocalCodingAgent] = {}
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self.allow_shell = allow_shell
        self.allow_write = allow_write

    def _account_base(self, account_id: str | None) -> Path:
        if not account_id:
            return self.session_directory.parent
        safe_id = _safe_account_id(account_id)
        return (self.session_directory.parent / 'accounts' / safe_id).resolve()

    def account_paths(self, account_id: str | None) -> dict[str, Path]:
        if not account_id:
            base = self.session_directory.parent
            return {
                'base': base,
                'sessions': self.session_directory,
                'scratchpad': self.session_directory,
                'uploads': self.session_directory,
                'outputs': self.session_directory,
            }
        base = self._account_base(account_id)
        return {
            'base': base,
            'sessions': base / 'sessions',
            'scratchpad': base / 'sessions',
            'uploads': base / 'sessions',
            'outputs': base / 'sessions',
        }

    def _build_agent(self, account_id: str | None = None) -> LocalCodingAgent:
        paths = self.account_paths(account_id)
        for directory in paths.values():
            directory.mkdir(parents=True, exist_ok=True)
        permissions = AgentPermissions(
            allow_file_write=self.allow_write,
            allow_shell_commands=self.allow_shell,
        )
        runtime_config = AgentRuntimeConfig(
            cwd=self.cwd,
            permissions=permissions,
            session_directory=paths['sessions'],
            scratchpad_root=paths['scratchpad'],
        )
        model_config = ModelConfig(
            model=self.model,
            base_url=self.base_url,
            api_key=self.api_key,
        )
        return LocalCodingAgent(
            model_config=model_config,
            runtime_config=runtime_config,
        )

    def agent_for(self, account_id: str | None = None) -> LocalCodingAgent:
        key = _safe_account_id(account_id) if account_id else '__default__'
        agent = self._agents.get(key)
        if agent is None:
            agent = self._build_agent(account_id)
            self._agents[key] = agent
        return agent

    def update(
        self,
        *,
        model: str | None = None,
        base_url: str | None = None,
        api_key: str | None = None,
        cwd: str | None = None,
        allow_shell: bool | None = None,
        allow_write: bool | None = None,
        account_id: str | None = None,
    ) -> None:
        with self._lock:
            # account_id 只用于让调用方拿到对应账号的 snapshot，不参与全局配置更新。
            del account_id
            if model is not None:
                self.model = _normalize_chat_model_name(model)
            if base_url is not None:
                self.base_url = base_url
            if api_key is not None:
                self.api_key = api_key
            if cwd is not None:
                resolved = Path(cwd).expanduser().resolve()
                if not resolved.is_dir():
                    raise ValueError(f'cwd does not exist: {resolved}')
                self.cwd = resolved
            if allow_shell is not None:
                self.allow_shell = allow_shell
            if allow_write is not None:
                self.allow_write = allow_write
            self._agents.clear()

    def snapshot(self, account_id: str | None = None) -> dict[str, Any]:
        agent = self.agent_for(account_id)
        paths = self.account_paths(account_id)
        return {
            'model': self.model,
            'base_url': self.base_url,
            'cwd': str(self.cwd),
            'account_id': _safe_account_id(account_id) if account_id else None,
            'account_directory': str(paths['base']),
            'session_directory': str(paths['sessions']),
            'upload_directory': str(paths['uploads']),
            'output_directory': str(paths['outputs']),
            'allow_shell': self.allow_shell,
            'allow_write': self.allow_write,
            'active_session_id': agent.active_session_id,
        }

    def lock(self) -> Lock:
        return self._lock


# ---------------------------------------------------------------------------
# Request models
# ---------------------------------------------------------------------------

class ChatRequest(BaseModel):
    prompt: str = Field(min_length=1)
    runtime_context: str | None = None
    resume_session_id: str | None = None
    account_id: str | None = None
    session_id: str | None = None


class StateUpdate(BaseModel):
    model: str | None = None
    base_url: str | None = None
    api_key: str | None = None
    cwd: str | None = None
    allow_shell: bool | None = None
    allow_write: bool | None = None
    account_id: str | None = None


class ModelListRequest(BaseModel):
    base_url: str | None = None
    api_key: str | None = None


# ---------------------------------------------------------------------------
# App factory
# ---------------------------------------------------------------------------

def create_app(state: AgentState) -> FastAPI:
    app = FastAPI(title='Claw Code GUI', version='1.0')

    # ------------- static + index ------------------------------------------
    app.mount(
        '/static',
        StaticFiles(directory=str(STATIC_DIR)),
        name='static',
    )

    @app.get('/', include_in_schema=False)
    async def root() -> FileResponse:
        return FileResponse(STATIC_DIR / 'index.html')

    # ------------- info ------------------------------------------------------
    @app.get('/api/state')
    async def get_state(account_id: str | None = None) -> dict[str, Any]:
        return state.snapshot(account_id)

    @app.post('/api/state')
    async def post_state(payload: StateUpdate) -> dict[str, Any]:
        try:
            state.update(**payload.model_dump(exclude_none=True))
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc))
        return state.snapshot(payload.account_id)

    @app.get('/api/slash-commands')
    async def list_slash_commands() -> list[dict[str, Any]]:
        commands: list[dict[str, Any]] = []
        for spec in get_slash_command_specs():
            commands.append(
                {
                    'names': list(spec.names),
                    'primary': spec.names[0],
                    'description': spec.description,
                }
            )
        return commands

    @app.get('/api/skills')
    async def list_skills() -> list[dict[str, Any]]:
        return [
            {
                'name': skill.name,
                'description': skill.description,
                'when_to_use': skill.when_to_use,
                'aliases': list(skill.aliases),
                'allowed_tools': list(skill.allowed_tools),
            }
            for skill in get_bundled_skills(state.cwd)
            if skill.user_invocable
        ]

    @app.get('/api/models')
    async def list_models_get() -> dict[str, Any]:
        return _list_backend_models(state.base_url, state.api_key)

    @app.post('/api/models')
    async def list_models_post(payload: ModelListRequest) -> dict[str, Any]:
        return _list_backend_models(
            payload.base_url or state.base_url,
            payload.api_key or state.api_key,
        )

    # ------------- sessions --------------------------------------------------
    @app.get('/api/sessions')
    async def list_sessions(account_id: str | None = None) -> list[dict[str, Any]]:
        directory = state.account_paths(account_id)['sessions']
        if not directory.exists():
            return []
        results: list[dict[str, Any]] = []
        for path in sorted(_iter_session_files(directory), key=_session_mtime, reverse=True):
            try:
                data = json.loads(path.read_text(encoding='utf-8'))
            except (OSError, json.JSONDecodeError):
                continue
            session_id = str(data.get('session_id', _session_id_from_path(path)))
            usage = data.get('usage') if isinstance(data.get('usage'), dict) else {}
            model_config = (
                data.get('model_config')
                if isinstance(data.get('model_config'), dict)
                else {}
            )
            messages = data.get('messages') or []
            preview = ''
            for msg in messages:
                if isinstance(msg, dict) and msg.get('role') == 'user':
                    content = msg.get('content', '')
                    if isinstance(content, str) and not _is_internal_message(content):
                        preview = _strip_session_context(content)[:120]
                        break
            results.append(
                {
                    'session_id': session_id,
                    'turns': data.get('turns', 0),
                    'tool_calls': data.get('tool_calls', 0),
                    'preview': preview,
                    'modified_at': _session_mtime(path),
                    'model': model_config.get('model'),
                    'usage': usage,
                }
            )
        return results

    @app.get('/api/sessions/{session_id}')
    async def get_session(session_id: str, account_id: str | None = None) -> dict[str, Any]:
        directory = state.account_paths(account_id)['sessions']
        try:
            stored = load_agent_session(session_id, directory=directory)
        except FileNotFoundError:
            raise HTTPException(status_code=404, detail='Session not found')
        return _serialize_stored_session(stored)

    # ------------- chat ------------------------------------------------------
    @app.post('/api/chat')
    async def chat(request: ChatRequest) -> dict[str, Any]:
        prompt = request.prompt.strip()
        if not prompt:
            raise HTTPException(status_code=400, detail='Prompt is empty')

        def _run() -> dict[str, Any]:
            with state.lock():
                agent = state.agent_for(request.account_id)
                session_directory = state.account_paths(request.account_id)['sessions']
                started_at = time.perf_counter()
                if request.resume_session_id is not None:
                    try:
                        stored = load_agent_session(
                            request.resume_session_id,
                            directory=session_directory,
                        )
                    except FileNotFoundError:
                        raise HTTPException(
                            status_code=404,
                            detail='Session to resume not found',
                        )
                    result = agent.resume(
                        prompt,
                        stored,
                        runtime_context=request.runtime_context,
                    )
                else:
                    result = agent.run(
                        prompt,
                        session_id=_safe_session_id(request.session_id),
                        runtime_context=request.runtime_context,
                    )
                elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000))
                payload = _serialize_run_result(result)
                payload['elapsed_ms'] = elapsed_ms
                if result.session_id:
                    _annotate_last_assistant_elapsed(
                        session_directory,
                        result.session_id,
                        elapsed_ms,
                    )
                    _annotate_transcript_elapsed(payload, elapsed_ms)
                return payload

        try:
            payload = await asyncio.to_thread(_run)
        except HTTPException:
            raise
        except Exception as exc:  # surface the error in the UI
            return JSONResponse(
                status_code=500,
                content={
                    'error': str(exc),
                    'error_type': type(exc).__name__,
                },
            )
        return payload

    @app.post('/api/clear')
    async def clear_state(account_id: str | None = None) -> dict[str, Any]:
        with state.lock():
            state.agent_for(account_id).clear_runtime_state()
        return state.snapshot(account_id)

    return app


# ---------------------------------------------------------------------------
# Serialization helpers
# ---------------------------------------------------------------------------

def _serialize_run_result(result: Any) -> dict[str, Any]:
    return {
        'final_output': result.final_output,
        'turns': result.turns,
        'tool_calls': result.tool_calls,
        'transcript': [_normalize_transcript_entry(entry) for entry in result.transcript],
        'session_id': result.session_id,
        'usage': result.usage.to_dict(),
        'total_cost_usd': result.total_cost_usd,
        'stop_reason': result.stop_reason,
    }


def _normalize_transcript_entry(entry: dict[str, Any]) -> dict[str, Any]:
    out: dict[str, Any] = {
        'role': entry.get('role', ''),
        'content': entry.get('content', ''),
    }
    for key in ('name', 'tool_call_id', 'tool_calls', 'metadata', 'message_id'):
        if key in entry and entry[key] not in (None, '', [], {}):
            out[key] = entry[key]
    return out


def _serialize_stored_session(stored: StoredAgentSession) -> dict[str, Any]:
    return {
        'session_id': stored.session_id,
        'turns': stored.turns,
        'tool_calls': stored.tool_calls,
        'messages': [_normalize_transcript_entry(dict(m)) for m in stored.messages],
        'usage': stored.usage,
        'total_cost_usd': stored.total_cost_usd,
        'model': stored.model_config.get('model'),
    }


def _list_backend_models(base_url: str, api_key: str) -> dict[str, Any]:
    req = request.Request(
        _join_url(base_url, '/models'),
        headers={
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json',
        },
        method='GET',
    )
    try:
        with request.urlopen(req, timeout=10) as response:
            payload = json.loads(response.read().decode('utf-8'))
    except error.HTTPError as exc:
        detail = exc.read().decode('utf-8', errors='replace')
        raise HTTPException(
            status_code=502,
            detail=f'HTTP {exc.code} from model backend: {detail}',
        ) from exc
    except (error.URLError, OSError, json.JSONDecodeError) as exc:
        raise HTTPException(
            status_code=502,
            detail=f'Unable to list models from {base_url}: {exc}',
        ) from exc

    models = _normalize_model_list(payload)
    return {'models': models, 'raw_count': len(models)}


def _normalize_model_list(payload: Any) -> list[dict[str, str]]:
    data = payload.get('data') if isinstance(payload, dict) else payload
    if not isinstance(data, list):
        return []
    model_ids: dict[str, str] = {}
    for item in data:
        if isinstance(item, str):
            normalized = _normalize_model_option(item)
            if normalized is not None:
                model_ids.setdefault(normalized.lower(), normalized)
            continue
        if not isinstance(item, dict):
            continue
        model_id = item.get('id') or item.get('model') or item.get('name')
        if not isinstance(model_id, str) or not model_id:
            continue
        normalized = _normalize_model_option(model_id)
        if normalized is not None:
            model_ids.setdefault(normalized.lower(), normalized)
    return [{'id': model_id} for model_id in sorted(model_ids.values(), key=str.lower)]


def _normalize_model_option(model_id: str) -> str | None:
    normalized = _normalize_chat_model_name(model_id)
    lower = normalized.lower()
    if lower.startswith('xiaomi/mimo-v2') and not any(
        blocked in lower for blocked in ('audio', 'asr', 'tts', 'voiceclone', 'voicedesign')
    ):
        return normalized
    known_chat_models = {
        'xiaomi/deepseek-r1-0528',
        'xiaomi/minimax-m2.5',
        'xiaomi/qwen3-32b',
    }
    if lower in known_chat_models:
        return normalized
    return None


def _normalize_chat_model_name(model: str) -> str:
    value = model.strip()
    lower = value.lower()
    if lower.startswith('mimo-v2') or lower in {
        'deepseek-r1-0528',
        'minimax-m2.5',
        'qwen3-32b',
    }:
        return f'xiaomi/{value}'
    return value


def _join_url(base_url: str, suffix: str) -> str:
    return f'{base_url.rstrip("/")}/{suffix.lstrip("/")}'


def _annotate_transcript_elapsed(payload: dict[str, Any], elapsed_ms: int) -> None:
    transcript = payload.get('transcript')
    if not isinstance(transcript, list):
        return
    for entry in reversed(transcript):
        if not isinstance(entry, dict) or entry.get('role') != 'assistant':
            continue
        metadata = entry.get('metadata')
        if not isinstance(metadata, dict):
            metadata = {}
            entry['metadata'] = metadata
        metadata['elapsed_ms'] = elapsed_ms
        return


def _annotate_last_assistant_elapsed(
    directory: Path,
    session_id: str,
    elapsed_ms: int,
) -> None:
    path = directory / session_id / 'session.json'
    if not path.exists():
        path = directory / f'{session_id}.json'
    try:
        data = json.loads(path.read_text(encoding='utf-8'))
    except (OSError, json.JSONDecodeError):
        return
    messages = data.get('messages')
    if not isinstance(messages, list):
        return
    for message in reversed(messages):
        if not isinstance(message, dict) or message.get('role') != 'assistant':
            continue
        metadata = message.get('metadata')
        if not isinstance(metadata, dict):
            metadata = {}
            message['metadata'] = metadata
        metadata['elapsed_ms'] = elapsed_ms
        try:
            path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
        except OSError:
            return
        return


def _safe_account_id(account_id: str | None) -> str:
    normalized = re.sub(r'[^a-zA-Z0-9._-]+', '_', (account_id or '').strip())
    return normalized[:80] or 'default'


def _is_internal_message(content: str) -> bool:
    return content.lstrip().startswith('<system-reminder>')


def _strip_session_context(content: str) -> str:
    marker = '\n\n[当前会话目录]'
    if marker in content:
        return content.split(marker, 1)[0].strip()
    marker = '\n[当前会话目录]'
    if marker in content:
        return content.split(marker, 1)[0].strip()
    return content.strip()


def _iter_session_files(directory: Path) -> list[Path]:
    if not directory.exists():
        return []
    by_session_id: dict[str, Path] = {}
    for path in directory.glob('*.json'):
        by_session_id[_session_id_from_path(path)] = path
    for path in directory.glob('*/session.json'):
        # 新目录格式包含 input/output/scratchpad，优先级高于旧的平铺 json。
        by_session_id[_session_id_from_path(path)] = path
    return list(by_session_id.values())


def _session_mtime(path: Path) -> float:
    try:
        return path.stat().st_mtime
    except OSError:
        return 0.0


def _session_id_from_path(path: Path) -> str:
    if path.name == 'session.json':
        return path.parent.name
    return path.stem


def _safe_session_id(session_id: str | None) -> str | None:
    if not session_id:
        return None
    normalized = re.sub(r'[^a-zA-Z0-9._-]+', '_', session_id.strip())
    return normalized[:120] or None