"""Collector: parses Claude Code agent output into structured Result objects."""

from __future__ import annotations

import base64
import contextlib
import json
import os
import re
from typing import Any

from dobby.models import Artifact, Result

# Regex to extract the DOBBY_RESULT block
_DOBBY_RESULT_RE = re.compile(
    r"---DOBBY_RESULT---\s*(.*?)\s*---END_DOBBY_RESULT---",
    re.DOTALL,
)

# Regex to find GitHub PR URLs
_PR_URL_RE = re.compile(
    r"https?://github\.com/[^/\s]+/[^/\s]+/pull/\d+"
)


class Collector:
    """Parses agent output into Result objects."""

    @staticmethod
    def _decode_codex_jsonl(output: str) -> str | None:
        """Try to decode Codex JSONL output.

        Codex outputs JSONL where each line is:
            {"id":"0","msg":{"type":"...","message":"..."}}
        Content chunks (type "agent" or "message") have base64-encoded message
        fields. Returns the concatenated decoded text, or None if the output
        is not Codex JSONL.
        """
        lines = output.strip().splitlines()
        if not lines:
            return None

        decoded_parts: list[str] = []
        valid_jsonl_lines = 0

        for line in lines:
            line = line.strip()
            if not line:
                continue
            try:
                event = json.loads(line)
            except (json.JSONDecodeError, TypeError):
                return None  # Not JSONL at all
            if not isinstance(event, dict) or "id" not in event or "msg" not in event:
                return None  # Not Codex JSONL format
            valid_jsonl_lines += 1
            msg = event["msg"]
            if not isinstance(msg, dict):
                continue
            content = msg.get("message", "")
            if not content:
                continue
            # Try base64 decoding; fall back to raw string
            try:
                decoded = base64.b64decode(content).decode("utf-8")
            except Exception:
                decoded = content
            decoded_parts.append(decoded)

        if valid_jsonl_lines == 0:
            return None

        return "".join(decoded_parts) if decoded_parts else ""

    @staticmethod
    def parse_output(output: str, exit_code: int, task_id: str) -> Result:
        """Parse agent output into a Result.

        Parsing order:
        1. Try to detect Codex JSONL format (multiline, each line is a JSON event
           with base64-encoded content). Decode and search for DOBBY_RESULT.
        2. Try to detect Claude Code JSON envelope (--output-format json), extract
           cost info and inner ``result`` field, then search that for DOBBY_RESULT.
        3. Try to extract a ``---DOBBY_RESULT---`` block from the raw output.
        4. Heuristic fallback: determine status from exit_code, scan for PR URLs.
        """
        cost: dict[str, Any] = {
            "input_tokens": 0,
            "output_tokens": 0,
            "total_tokens": 0,
            "cost_usd": 0.0,
        }

        # 1. Try Codex JSONL format
        codex_text = Collector._decode_codex_jsonl(output)
        if codex_text is not None:
            match = _DOBBY_RESULT_RE.search(codex_text)
            if match:
                try:
                    data = json.loads(match.group(1))
                    artifacts: list[Artifact] = []
                    pr_url = data.get("pr_url")
                    if pr_url:
                        artifacts.append(Artifact(type="pr", url=pr_url))
                    return Result(
                        task_id=task_id,
                        status=data.get("status", "completed"),
                        summary=data.get("summary", ""),
                        artifacts=artifacts,
                        metrics=data.get("metrics") or {},
                        cost=cost,
                    )
                except (json.JSONDecodeError, TypeError):
                    pass
            # Codex JSONL detected but no DOBBY_RESULT — use decoded text for heuristic
            return Collector._heuristic_parse(codex_text, exit_code, task_id, cost)

        # 2. Try Claude Code stream-json or single-line JSON
        inner_text: str | None = None
        session_id: str | None = None
        lines = output.strip().splitlines()

        # stream-json: multi-line JSONL, last line has type=result
        if len(lines) > 1:
            for line in reversed(lines):
                line = line.strip()
                if not line:
                    continue
                try:
                    event = json.loads(line)
                    if isinstance(event, dict) and event.get("type") == "result":
                        # Extract cost from stream-json result event
                        if "total_cost_usd" in event:
                            cost["cost_usd"] = event["total_cost_usd"]
                        usage = event.get("usage", {})
                        if usage:
                            cost["input_tokens"] = (
                                usage.get("input_tokens", 0)
                                + usage.get("cache_read_input_tokens", 0)
                            )
                            cost["output_tokens"] = usage.get("output_tokens", 0)
                            cost["total_tokens"] = cost["input_tokens"] + cost["output_tokens"]
                        inner_text = event.get("result", "")
                        session_id = event.get("session_id")
                        break
                except (json.JSONDecodeError, TypeError):
                    continue

        # Single-line JSON envelope (legacy --output-format json)
        if inner_text is None and len(lines) == 1 and output.strip().startswith("{"):
            try:
                envelope = json.loads(output)
                if isinstance(envelope, dict) and "result" in envelope:
                    if "cost_usd" in envelope:
                        cost["cost_usd"] = envelope["cost_usd"]
                    if "input_tokens" in envelope:
                        cost["input_tokens"] = envelope["input_tokens"]
                    if "output_tokens" in envelope:
                        cost["output_tokens"] = envelope["output_tokens"]
                    cost["total_tokens"] = (
                        cost["input_tokens"] + cost["output_tokens"]
                    )
                    inner_text = envelope["result"]
                    if session_id is None:
                        session_id = envelope.get("session_id")
            except (json.JSONDecodeError, TypeError):
                pass

        search_text = inner_text if inner_text is not None else output

        # 3. Try DOBBY_RESULT block
        match = _DOBBY_RESULT_RE.search(search_text)
        if match:
            try:
                data = json.loads(match.group(1))
                artifacts: list[Artifact] = []
                pr_url = data.get("pr_url")
                if pr_url:
                    artifacts.append(Artifact(type="pr", url=pr_url))
                return Result(
                    task_id=task_id,
                    status=data.get("status", "completed"),
                    summary=data.get("summary", ""),
                    artifacts=artifacts,
                    metrics=data.get("metrics") or {},
                    cost=cost,
                    session_id=session_id,
                )
            except (json.JSONDecodeError, TypeError):
                pass

        # 4. Heuristic fallback
        result = Collector._heuristic_parse(output, exit_code, task_id, cost)
        result.session_id = session_id
        return result

    @staticmethod
    def _heuristic_parse(
        output: str,
        exit_code: int,
        task_id: str,
        cost: dict[str, Any],
    ) -> Result:
        """Fallback heuristic parser."""
        # Handle empty output
        if not output.strip():
            return Result(
                task_id=task_id,
                status="failed",
                summary=f"Agent produced no output and exited with code {exit_code}",
                cost=cost,
            )

        # Determine status from exit code
        status = "completed" if exit_code == 0 else "failed"

        # Scan for PR URLs
        artifacts: list[Artifact] = []
        pr_matches = _PR_URL_RE.findall(output)
        for url in pr_matches:
            artifacts.append(Artifact(type="pr", url=url))

        # Build summary from output — find the LAST assistant text (conclusion)
        summary = ""
        metrics: dict[str, Any] = {}
        lines = output.strip().splitlines()

        # Collect ALL assistant text blocks (last one is usually the conclusion)
        assistant_texts: list[str] = []
        for line in lines:
            line = line.strip()
            if not line or not line.startswith("{"):
                continue
            try:
                obj = json.loads(line)
                if not isinstance(obj, dict):
                    continue
                # Skip system init, user messages, result events
                if obj.get("type") in ("system", "user", "result"):
                    continue
                if obj.get("type") == "assistant":
                    content = obj.get("message", {}).get("content", [])
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            text = block["text"].strip()
                            if text and len(text) > 20:
                                assistant_texts.append(text)
            except (json.JSONDecodeError, TypeError):
                # Non-JSON line — collect as potential summary
                if not summary and len(line) > 20:
                    summary = line[:300]

        # Use the LAST assistant text as summary (it's usually the conclusion)
        if assistant_texts:
            summary = assistant_texts[-1][:2000]
            # Also try to extract metrics from the last text
            # Look for JSON-like structures in the text
            json_match = re.search(r'\{[^{}]*"accuracy"[^{}]*\}', assistant_texts[-1])
            if json_match:
                with contextlib.suppress(json.JSONDecodeError, TypeError):
                    metrics = json.loads(json_match.group())

        if not summary:
            summary = f"Task {status} (exit code {exit_code})"

        return Result(
            task_id=task_id,
            status=status,
            summary=summary,
            artifacts=artifacts,
            metrics=metrics,
            cost=cost,
        )

    @staticmethod
    def save_artifacts(task_id: str, output: str, results_dir: str) -> str:
        """Save raw agent log to results_dir/task_id/agent.log.

        Returns the path to the saved log file.
        """
        task_dir = os.path.join(results_dir, task_id)
        os.makedirs(task_dir, exist_ok=True)
        log_path = os.path.join(task_dir, "agent.log")
        with open(log_path, "w", encoding="utf-8") as fh:
            fh.write(output)
        return log_path
