import json
import os
import shlex
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from litellm.utils import get_model_info

from harbor.agents.installed.base import BaseInstalledAgent, ExecInput
from harbor.models.agent.context import AgentContext
from harbor.models.agent.name import AgentName
from harbor.models.trajectories import (
    Agent,
    FinalMetrics,
    Metrics,
    Observation,
    ObservationResult,
    Step,
    ToolCall,
    Trajectory,
)
from harbor.utils.trajectory_utils import format_trajectory_json

_PROVIDER_CONFIG: dict[str, dict[str, Any]] = {
    "moonshot": {
        "type": "kimi",
        "base_url": "https://api.moonshot.cn/v1",
        "env_keys": ["MOONSHOT_API_KEY"],
    },
    "kimi": {
        "type": "kimi",
        "base_url": "https://api.kimi.com/coding/v1",
        "env_keys": ["KIMI_API_KEY", "MOONSHOT_API_KEY"],
    },
    "openai": {
        "type": "openai_legacy",
        "base_url": "https://api.openai.com/v1",
        "env_keys": ["OPENAI_API_KEY"],
    },
    "anthropic": {
        "type": "anthropic",
        "base_url": "https://api.anthropic.com",
        "env_keys": ["ANTHROPIC_API_KEY"],
    },
    "gemini": {
        "type": "gemini",
        "base_url": "https://generativelanguage.googleapis.com",
        "env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY"],
    },
    "google": {
        "type": "gemini",
        "base_url": "https://generativelanguage.googleapis.com",
        "env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY"],
    },
}

_OUTPUT_FILENAME = "kimi-cli.txt"


@dataclass
class _PendingToolCall:
    call_id: str
    name: str
    arguments_buffer: str


@dataclass
class _WireStep:
    n: int
    text_parts: list[str] = field(default_factory=list)
    reasoning_parts: list[str] = field(default_factory=list)
    tool_calls: list[dict[str, Any]] = field(default_factory=list)
    tool_results: dict[str, dict[str, Any]] = field(default_factory=dict)
    token_usage: dict[str, Any] | None = None
    pending_tool: _PendingToolCall | None = None

    def finalize_pending_tool(self) -> None:
        if self.pending_tool is None:
            return
        buf = self.pending_tool.arguments_buffer
        try:
            arguments = json.loads(buf) if buf else {}
        except json.JSONDecodeError:
            arguments = {"raw": buf} if buf else {}
        if not isinstance(arguments, dict):
            arguments = {"value": arguments}
        self.tool_calls.append(
            {
                "id": self.pending_tool.call_id,
                "name": self.pending_tool.name,
                "arguments": arguments,
            }
        )
        self.pending_tool = None


class KimiCli(BaseInstalledAgent):
    def get_version_command(self) -> str | None:
        return "kimi --version"

    SUPPORTS_ATIF: bool = True

    _DEFAULT_MAX_CONTEXT_SIZE: int = 131072

    def __init__(self, *args, **kwargs):
        self._api_key: str | None = kwargs.pop("api_key", None)
        self._base_url: str | None = kwargs.pop("base_url", None)
        model_info: dict[str, Any] | None = kwargs.pop("model_info", None)
        super().__init__(*args, **kwargs)

        if model_info and self.model_name:
            import litellm

            litellm.register_model({self.model_name: model_info})

        self._max_context_size = self._resolve_max_context_size()

    @staticmethod
    def name() -> str:
        return AgentName.KIMI_CLI.value

    @property
    def _install_agent_template_path(self) -> Path:
        return Path(__file__).parent / "install-kimi-cli.sh.j2"

    def _resolve_api_key(self, provider: str) -> str:
        if self._api_key:
            return self._api_key
        pcfg = _PROVIDER_CONFIG.get(provider, {})
        for key in pcfg.get("env_keys", []):
            val = os.environ.get(key)
            if val:
                return val
        return ""

    def _resolve_max_context_size(self) -> int:
        if not self.model_name:
            return self._DEFAULT_MAX_CONTEXT_SIZE
        try:
            info = get_model_info(self.model_name)
            return info.get("max_input_tokens") or self._DEFAULT_MAX_CONTEXT_SIZE
        except Exception:
            return self._DEFAULT_MAX_CONTEXT_SIZE

    def _build_config_json(self, provider: str, model: str) -> str:
        pcfg = _PROVIDER_CONFIG.get(provider)
        if pcfg is None:
            raise ValueError(
                f"Unsupported provider '{provider}' for kimi-cli. "
                f"Supported: {sorted(_PROVIDER_CONFIG)}"
            )
        base_url = self._base_url or pcfg["base_url"]
        api_key = self._resolve_api_key(provider)
        config: dict[str, Any] = {
            "default_model": "model",
            "default_yolo": True,
            "providers": {
                "harbor": {
                    "type": pcfg["type"],
                    "base_url": base_url,
                    "api_key": api_key,
                }
            },
            "models": {
                "model": {
                    "provider": "harbor",
                    "model": model,
                    "max_context_size": self._max_context_size,
                }
            },
        }
        return json.dumps(config)

    def _build_mcp_config_json(self) -> str | None:
        if not self.mcp_servers:
            return None
        servers: dict[str, dict[str, Any]] = {}
        for server in self.mcp_servers:
            if server.transport == "stdio":
                entry: dict[str, Any] = {"command": server.command, "args": server.args}
            else:
                entry = {"url": server.url}
            servers[server.name] = entry
        return json.dumps({"mcpServers": servers})

    def populate_context_post_run(self, context: AgentContext) -> None:
        events = self._parse_wire_events()
        if not events:
            return
        try:
            trajectory = self._convert_events_to_trajectory(events)
        except Exception:
            self.logger.exception("Failed to convert kimi-cli events to trajectory")
            return
        if not trajectory:
            return

        trajectory_path = self.logs_dir / "trajectory.json"
        try:
            trajectory_path.write_text(
                format_trajectory_json(trajectory.to_json_dict())
            )
        except OSError as exc:
            print(f"Failed to write trajectory file {trajectory_path}: {exc}")

        if trajectory.final_metrics:
            fm = trajectory.final_metrics
            context.cost_usd = fm.total_cost_usd
            context.n_input_tokens = fm.total_prompt_tokens or 0
            context.n_output_tokens = fm.total_completion_tokens or 0
            context.n_cache_tokens = fm.total_cached_tokens or 0

    def _build_register_skills_command(self) -> str | None:
        if not self.skills_dir:
            return None
        return (
            f"mkdir -p ~/.kimi/skills && "
            f"cp -r {shlex.quote(self.skills_dir)}/* "
            f"~/.kimi/skills/ 2>/dev/null || true"
        )

    def _build_register_mcp_servers_command(self) -> str | None:
        mcp_json = self._build_mcp_config_json()
        if not mcp_json:
            return None
        escaped = shlex.quote(mcp_json)
        return f"mkdir -p /tmp && echo {escaped} > /tmp/kimi-mcp.json"

    def create_run_agent_commands(self, instruction: str) -> list[ExecInput]:
        if not self.model_name or "/" not in self.model_name:
            raise ValueError("Model name must be in format provider/model_name")

        provider, model = self.model_name.split("/", 1)

        config_json = self._build_config_json(provider, model)
        escaped_config = shlex.quote(config_json)

        prompt_request = json.dumps(
            {
                "jsonrpc": "2.0",
                "method": "prompt",
                "id": "1",
                "params": {"user_input": instruction},
            }
        )
        escaped_prompt = shlex.quote(prompt_request)

        env: dict[str, str] = {}
        pcfg = _PROVIDER_CONFIG.get(provider, {})
        for key in pcfg.get("env_keys", []):
            val = os.environ.get(key)
            if val:
                env[key] = val

        commands: list[ExecInput] = []

        setup_parts = [
            f"echo {escaped_config} > /tmp/kimi-config.json",
        ]

        skills_cmd = self._build_register_skills_command()
        if skills_cmd:
            setup_parts.append(skills_cmd)

        mcp_cmd = self._build_register_mcp_servers_command()
        if mcp_cmd:
            setup_parts.append(mcp_cmd)

        commands.append(ExecInput(command=" && ".join(setup_parts), env=env))

        mcp_flag = "--mcp-config-file /tmp/kimi-mcp.json " if mcp_cmd else ""

        run_command = (
            f'export PATH="$HOME/.local/bin:$PATH"; '
            f"(echo {escaped_prompt}; sleep 86400) | "
            f"kimi --config-file /tmp/kimi-config.json --wire --yolo "
            f"{mcp_flag}"
            f"2>/dev/null | ("
            f"while IFS= read -r line; do "
            f'echo "$line" >> /logs/agent/{_OUTPUT_FILENAME}; '
            'case "$line" in *\'"id":"1"\'*) break ;; esac; '
            f"done; kill 0 2>/dev/null)"
        )

        commands.append(ExecInput(command=run_command, env=env))
        return commands

    def _parse_wire_events(self) -> list[dict[str, Any]]:
        """Parse wire protocol JSONL, handling unescaped control characters
        in ToolResult output that may split a message across multiple lines."""
        output_path = self.logs_dir / _OUTPUT_FILENAME
        if not output_path.exists():
            return []
        events: list[dict[str, Any]] = []
        buffer = ""
        for line in output_path.read_text().splitlines():
            if line.lstrip().startswith('{"jsonrpc"'):
                if buffer:
                    self._try_parse_event(buffer, events)
                buffer = line
            elif buffer:
                buffer += "\n" + line
        if buffer:
            self._try_parse_event(buffer, events)
        return events

    @staticmethod
    def _try_parse_event(raw: str, out: list[dict[str, Any]]) -> None:
        try:
            msg = json.loads(raw, strict=False)
        except json.JSONDecodeError:
            return
        if msg.get("method") == "event":
            out.append(msg.get("params", {}))

    @staticmethod
    def _group_events_into_steps(
        events: list[dict[str, Any]],
    ) -> list[_WireStep]:
        steps: list[_WireStep] = []
        current: _WireStep | None = None

        for event in events:
            etype = event.get("type")
            payload = event.get("payload", {})

            if etype == "StepBegin":
                if current is not None:
                    current.finalize_pending_tool()
                    steps.append(current)
                current = _WireStep(n=payload.get("n", len(steps) + 1))
                continue

            if current is None:
                continue

            if etype == "ContentPart":
                content_type = payload.get("type")
                if content_type == "text":
                    text = payload.get("text", "")
                    if text:
                        current.text_parts.append(text)
                elif content_type == "think":
                    think = payload.get("think", "")
                    if think:
                        current.reasoning_parts.append(think)

            elif etype == "ToolCall":
                current.finalize_pending_tool()
                func = payload.get("function", {})
                current.pending_tool = _PendingToolCall(
                    call_id=payload.get("id", ""),
                    name=func.get("name", ""),
                    arguments_buffer=func.get("arguments", "") or "",
                )

            elif etype == "ToolCallPart":
                if current.pending_tool is not None:
                    current.pending_tool.arguments_buffer += (
                        payload.get("arguments_part") or ""
                    )

            elif etype == "ToolResult":
                current.finalize_pending_tool()
                tc_id = payload.get("tool_call_id", "")
                ret = payload.get("return_value", {})
                current.tool_results[tc_id] = ret

            elif etype == "StatusUpdate":
                current.token_usage = payload.get("token_usage")

            elif etype == "TurnEnd":
                current.finalize_pending_tool()
                steps.append(current)
                current = None

        if current is not None:
            current.finalize_pending_tool()
            steps.append(current)

        return steps

    def _convert_events_to_trajectory(
        self, events: list[dict[str, Any]]
    ) -> Trajectory | None:
        wire_steps = self._group_events_into_steps(events)
        if not wire_steps:
            return None

        steps: list[Step] = []
        total_prompt = 0
        total_completion = 0
        total_cached = 0

        for idx, ws in enumerate(wire_steps, start=1):
            message = "".join(ws.text_parts) or "(tool use)"
            reasoning = "".join(ws.reasoning_parts) or None

            tool_calls: list[ToolCall] | None = None
            observation: Observation | None = None

            if ws.tool_calls:
                tool_calls = []
                obs_results: list[ObservationResult] = []
                for tc in ws.tool_calls:
                    tc_id = tc["id"]
                    tool_calls.append(
                        ToolCall(
                            tool_call_id=tc_id,
                            function_name=tc["name"],
                            arguments=tc["arguments"],
                        )
                    )
                    result = ws.tool_results.get(tc_id)
                    if result is not None:
                        output = result.get("output")
                        if isinstance(output, list):
                            output = json.dumps(output, ensure_ascii=False)
                        obs_results.append(
                            ObservationResult(
                                source_call_id=tc_id,
                                content=str(output) if output is not None else None,
                            )
                        )
                if obs_results:
                    observation = Observation(results=obs_results)

            metrics: Metrics | None = None
            if ws.token_usage:
                tu = ws.token_usage
                input_other = tu.get("input_other", 0) or 0
                output_tok = tu.get("output", 0) or 0
                cache_read = tu.get("input_cache_read", 0) or 0
                cache_creation = tu.get("input_cache_creation", 0) or 0

                prompt_tokens = input_other + cache_read + cache_creation
                total_prompt += prompt_tokens
                total_completion += output_tok
                total_cached += cache_read

                extra: dict[str, Any] = {}
                if cache_creation:
                    extra["input_cache_creation"] = cache_creation
                metrics = Metrics(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=output_tok,
                    cached_tokens=cache_read if cache_read else None,
                    extra=extra or None,
                )

            step_kwargs: dict[str, Any] = {
                "step_id": idx,
                "source": "agent",
                "message": message,
                "model_name": self.model_name,
            }
            if tool_calls:
                step_kwargs["tool_calls"] = tool_calls
            if observation:
                step_kwargs["observation"] = observation
            if metrics:
                step_kwargs["metrics"] = metrics
            if reasoning:
                step_kwargs["reasoning_content"] = reasoning

            steps.append(Step(**step_kwargs))

        if not steps:
            return None

        final_metrics = FinalMetrics(
            total_prompt_tokens=total_prompt or None,
            total_completion_tokens=total_completion or None,
            total_cached_tokens=total_cached or None,
            total_cost_usd=None,
            total_steps=len(steps),
        )

        return Trajectory(
            schema_version="ATIF-v1.6",
            session_id="unknown",
            agent=Agent(
                name="kimi-cli",
                version=self.version() or "unknown",
                model_name=self.model_name,
            ),
            steps=steps,
            final_metrics=final_metrics,
        )
