import json
import os
import shlex
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.models.trial.paths import EnvironmentPaths


class Codex(BaseInstalledAgent):
    """
    The Codex agent uses OpenAI's Codex CLI tool to solve tasks.
    """

    _OUTPUT_FILENAME = "codex.txt"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._version = kwargs.get("version", None)

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

    @property
    def _trajectory_path(self) -> Path:
        return EnvironmentPaths.agent_dir / "trajectory.json"

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

    def _get_session_dir(self) -> Path | None:
        """Get the single session directory."""
        sessions_dir = self.logs_dir / "sessions"
        if not sessions_dir.exists():
            return None

        session_dirs = [d for d in sessions_dir.rglob("*") if d.is_dir()]
        session_dirs = [
            d
            for d in session_dirs
            if len(d.parts) == max(len(x.parts) for x in session_dirs)
        ]
        if not session_dirs:
            return None

        # Sanity check: there should be exactly one session
        assert len(session_dirs) == 1, (
            f"Expected exactly 1 session, found {len(session_dirs)}"
        )
        return session_dirs[0]

    @staticmethod
    def _extract_message_text(content: list[Any]) -> str:
        """Extract joined text from Codex content blocks."""
        parts: list[str] = []
        for block in content:
            if isinstance(block, dict):
                text = block.get("text")
                if isinstance(text, str):
                    parts.append(text)
        return "".join(parts)

    @staticmethod
    def _parse_output_blob(raw: Any) -> tuple[str | None, dict[str, Any] | None]:
        """Extract textual output and metadata from Codex tool outputs."""
        if raw is None:
            return None, None

        if isinstance(raw, str):
            try:
                parsed = json.loads(raw)
            except json.JSONDecodeError:
                return raw, None
        else:
            parsed = raw

        if isinstance(parsed, dict):
            output = parsed.get("output")
            if output is None and parsed:
                # dumping remaining structure if output missing
                output = json.dumps(parsed, ensure_ascii=False)
            metadata = parsed.get("metadata")
            return output, metadata if isinstance(metadata, dict) else None

        return str(parsed), None

    def _convert_event_to_step(self, event: dict[str, Any], step_id: int) -> Step:
        """Convert a normalized Codex event dictionary into an ATIF step."""
        kind = event.get("kind")
        timestamp = event.get("timestamp")

        if kind == "message":
            role = event.get("role", "user")
            text = event.get("text", "")
            reasoning = event.get("reasoning")
            source: str
            if role == "assistant":
                source = "agent"
            elif role == "user":
                source = "user"
            else:
                source = "system"

            kwargs: dict[str, Any] = {
                "step_id": step_id,
                "timestamp": timestamp,
                "source": source,
                "message": text,
            }

            if source == "agent":
                if reasoning:
                    kwargs["reasoning_content"] = reasoning
                if self.model_name:
                    kwargs["model_name"] = self.model_name

            extra = event.get("extra")
            if extra:
                kwargs["extra"] = extra

            return Step(**kwargs)

        if kind == "tool_call":
            call_id = event.get("call_id", "")
            tool_name = event.get("tool_name", "")
            reasoning = event.get("reasoning")
            arguments = event.get("arguments") or {}
            if not isinstance(arguments, dict):
                arguments = {"value": arguments}

            tool_call = ToolCall(
                tool_call_id=call_id,
                function_name=tool_name,
                arguments=arguments,
            )

            observation: Observation | None = None
            output_text = event.get("output")
            if output_text is not None:
                observation = Observation(
                    results=[
                        ObservationResult(
                            source_call_id=call_id or None,
                            content=output_text,
                        )
                    ]
                )

            metrics_payload = event.get("metrics")
            metrics: Metrics | None = None
            if isinstance(metrics_payload, dict):
                metrics = Metrics(**metrics_payload)

            extra: dict[str, Any] | None = None
            metadata = event.get("metadata")
            if metadata:
                extra = {"tool_metadata": metadata}
            raw_arguments = event.get("raw_arguments")
            if raw_arguments:
                extra = extra or {}
                extra["raw_arguments"] = raw_arguments
            status = event.get("status")
            if status:
                extra = extra or {}
                extra["status"] = status

            message = event.get("message")
            if not message:
                summary_parts = [part for part in [tool_name, call_id] if part]
                summary = " ".join(summary_parts) if summary_parts else "Tool call"
                message = f"Executed {summary}"

            kwargs = {
                "step_id": step_id,
                "timestamp": timestamp,
                "source": "agent",
                "message": message,
                "tool_calls": [tool_call],
                "observation": observation,
            }

            if self.model_name:
                kwargs["model_name"] = self.model_name
            if reasoning:
                kwargs["reasoning_content"] = reasoning
            if metrics:
                kwargs["metrics"] = metrics
            if extra:
                kwargs["extra"] = extra

            return Step(**kwargs)

        raise ValueError(f"Unsupported event kind '{kind}'")

    def _convert_events_to_trajectory(self, session_dir: Path) -> Trajectory | None:
        """Convert Codex session JSONL events into an ATIF trajectory."""
        session_files = list(session_dir.glob("*.jsonl"))

        if not session_files:
            print(f"No Codex session files found in {session_dir}")
            return None

        session_file = session_files[0]
        if not session_file:
            print(f"No Codex session files found in {session_dir}")
            return None

        raw_events: list[dict[str, Any]] = []
        with open(session_file, "r") as handle:
            for line in handle:
                stripped = line.strip()
                if not stripped:
                    continue
                try:
                    raw_events.append(json.loads(stripped))
                except json.JSONDecodeError as exc:
                    print(f"Skipping malformed JSONL line in {session_file}: {exc}")

        if not raw_events:
            return None

        session_meta = next(
            (e for e in raw_events if e.get("type") == "session_meta"), None
        )
        session_id = (
            session_meta.get("payload", {}).get("id")
            if session_meta and isinstance(session_meta, dict)
            else session_dir.name
        )

        agent_version = "unknown"
        agent_extra: dict[str, Any] | None = None
        default_model_name: str | None = None

        if session_meta:
            payload = session_meta.get("payload", {})
            agent_version = payload.get("cli_version") or agent_version
            extra: dict[str, Any] = {}
            for key in ("originator", "cwd", "git", "instructions"):
                value = payload.get(key)
                if value is not None:
                    extra[key] = value
            agent_extra = extra or None

        if agent_version == "unknown" and self._version:
            agent_version = self._version

        for event in raw_events:
            if event.get("type") == "turn_context":
                model_name = event.get("payload", {}).get("model")
                if isinstance(model_name, str):
                    default_model_name = model_name
                    break

        if default_model_name is None:
            default_model_name = self.model_name

        # normalize events to a structure suitable for conversion into Steps
        normalized_events: list[dict[str, Any]] = []
        pending_calls: dict[str, dict[str, Any]] = {}
        pending_reasoning: str | None = None

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

            if etype != "response_item":
                if etype == "event_msg":
                    # token counts handled later for metrics
                    continue
                continue

            payload_type = payload.get("type")
            if payload_type == "reasoning":
                summary = payload.get("summary")
                if isinstance(summary, list) and summary:
                    pending_reasoning = "\n".join(
                        str(item) for item in summary if isinstance(item, str)
                    )
                else:
                    pending_reasoning = None
                continue

            if payload_type == "message":
                content = payload.get("content", [])
                text = (
                    self._extract_message_text(content)
                    if isinstance(content, list)
                    else ""
                )
                normalized_events.append(
                    {
                        "kind": "message",
                        "timestamp": timestamp,
                        "role": payload.get("role", "user"),
                        "text": text,
                        "reasoning": pending_reasoning
                        if payload.get("role") == "assistant"
                        else None,
                    }
                )
                pending_reasoning = None
                continue

            if payload_type in {"function_call", "custom_tool_call"}:
                call_id = payload.get("call_id")
                if not call_id:
                    continue

                raw_args_key = (
                    "arguments" if payload_type == "function_call" else "input"
                )
                raw_arguments = payload.get(raw_args_key)
                try:
                    parsed_args = json.loads(raw_arguments)
                except (json.JSONDecodeError, TypeError):
                    raise ValueError(f"Malformed tool arguments for call_id {call_id}")

                pending_calls[call_id] = {
                    "kind": "tool_call",
                    "timestamp": timestamp,
                    "call_id": call_id,
                    "tool_name": payload.get("name") or "",
                    "arguments": parsed_args,
                    "raw_arguments": raw_arguments,
                    "reasoning": pending_reasoning,
                    "status": payload.get("status"),
                    "message": None,
                }
                pending_reasoning = None
                continue

            if payload_type in {"function_call_output", "custom_tool_call_output"}:
                call_id = payload.get("call_id")
                output_text, metadata = self._parse_output_blob(payload.get("output"))

                call_info = pending_calls.pop(call_id, None) if call_id else None

                if call_info is None:
                    call_info = {
                        "kind": "tool_call",
                        "timestamp": timestamp,
                        "call_id": call_id or "",
                        "tool_name": payload.get("name", "") or "",
                        "arguments": {},
                        "raw_arguments": None,
                        "reasoning": pending_reasoning,
                        "status": None,
                        "message": None,
                    }

                call_info["output"] = output_text
                call_info["metadata"] = metadata
                call_info["timestamp"] = call_info.get("timestamp") or timestamp
                normalized_events.append(call_info)
                pending_reasoning = None
                continue

        steps: list[Step] = []
        for idx, norm_event in enumerate(normalized_events, start=1):
            try:
                step = self._convert_event_to_step(norm_event, idx)
            except ValueError as exc:
                print(f"Skipping event during step conversion: {exc}")
                continue

            # Provide default model name if not set for agent steps
            if step.source == "agent" and not step.model_name and default_model_name:
                step.model_name = default_model_name

            steps.append(step)

        if not steps:
            print("No valid steps produced from Codex session")
            return None

        # Extract final metrics from the last token_count event with totals
        total_metrics: FinalMetrics | None = None
        for event in reversed(raw_events):
            if event.get("type") != "event_msg":
                continue
            payload = event.get("payload", {})
            if not isinstance(payload, dict) or payload.get("type") != "token_count":
                continue

            info = payload.get("info")
            if not isinstance(info, dict):
                continue

            total_usage = info.get("total_token_usage")
            if not isinstance(total_usage, dict):
                continue

            prompt_tokens = total_usage.get("input_tokens")
            completion_tokens = total_usage.get("output_tokens")
            reasoning_tokens = total_usage.get("reasoning_output_tokens")
            cached_tokens = total_usage.get("cached_input_tokens")
            overall_tokens = total_usage.get("total_tokens")

            total_metrics = FinalMetrics(
                total_prompt_tokens=prompt_tokens if prompt_tokens else None,
                total_completion_tokens=completion_tokens or None,
                total_cached_tokens=cached_tokens or None,
                total_cost_usd=info.get("total_cost") or info.get("cost_usd"),
                total_steps=len(steps),
                extra={
                    "reasoning_output_tokens": reasoning_tokens,
                    "total_tokens": overall_tokens,
                    "last_token_usage": info.get("last_token_usage"),
                },
            )
            break

        trajectory = Trajectory(
            schema_version="ATIF-v1.2",
            session_id=session_id,
            agent=Agent(
                name="codex",
                version=agent_version,
                model_name=default_model_name,
                extra=agent_extra,
            ),
            steps=steps,
            final_metrics=total_metrics,
        )

        return trajectory

    def populate_context_post_run(self, context: AgentContext) -> None:
        """
        Populate the agent context after Codex finishes executing.

        Converts the Codex session JSONL file into an ATIF trajectory, persists it,
        and propagates usage metrics back to the Harbor context.
        """
        session_dir = self._get_session_dir()
        if not session_dir:
            print("No Codex session directory found")
            return

        try:
            trajectory = self._convert_events_to_trajectory(session_dir)
        except Exception as exc:
            print(f"Failed to convert Codex events to trajectory: {exc}")
            return

        if not trajectory:
            print("Failed to convert Codex session to trajectory")
            return

        trajectory_path = self.logs_dir / "trajectory.json"
        try:
            with open(trajectory_path, "w") as handle:
                json.dump(trajectory.to_json_dict(), handle, indent=2)
            print(f"Wrote Codex trajectory to {trajectory_path}")
        except OSError as exc:
            print(f"Failed to write trajectory file {trajectory_path}: {exc}")

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

    def create_run_agent_commands(self, instruction: str) -> list[ExecInput]:
        escaped_instruction = shlex.quote(instruction)

        if not self.model_name or "/" not in self.model_name:
            raise ValueError("Model name must be in the format provider/model_name")

        model = self.model_name.split("/")[-1]

        env = {
            "OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", ""),
            "CODEX_HOME": (EnvironmentPaths.agent_dir).as_posix(),
        }

        return [
            ExecInput(
                command="""
cat >"$CODEX_HOME/auth.json" <<EOF
{
  "OPENAI_API_KEY": "${OPENAI_API_KEY}"
}
EOF
                """,
                env=env,
            ),
            ExecInput(
                command=(
                    "codex exec "
                    "--dangerously-bypass-approvals-and-sandbox "
                    "--skip-git-repo-check "
                    f"--model {model} "
                    "--json "
                    "-- "  # end of flags
                    f"{escaped_instruction} "
                    f"2>&1 </dev/null | tee {
                        EnvironmentPaths.agent_dir / self._OUTPUT_FILENAME
                    } "
                    "&& rm -rf $CODEX_HOME/auth.json"
                ),
                env=env,
            ),
        ]
