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

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

    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)
        super().__init__(*args, **kwargs)

    @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 _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": 131072,
                }
            },
        }
        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,
        )
