"""Enhanced CLI adapter with structured output and accurate cost tracking."""

import asyncio
import json
import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

from deliberate.adapters.base import AdapterResponse, ModelAdapter
from deliberate.settings import DeliberateSettings, get_agent_env_dict

try:
    from genai_prices import get_model_price
    HAS_GENAI_PRICES = True
except ImportError:
    HAS_GENAI_PRICES = False


@dataclass
class CLIAdapter(ModelAdapter):
    """CLI adapter with native structured output support.

    Intelligently adds CLI-native flags for:
    - JSON schema validation (claude --json-schema, codex --output-schema)
    - JSON output format (all CLIs support)
    - Headless operation (permissions bypassing)

    Integrates genai-prices for accurate cost calculation when available.
    """

    name: str
    command: list[str]
    model_id: str | None = None  # For genai-prices lookup
    env: dict[str, str] = field(default_factory=dict)
    timeout_seconds: int = 300
    _call_count: int = field(default=0, repr=False)

    def __post_init__(self):
        """Auto-detect CLI tool and model if not specified."""
        if not self.model_id and self.command:
            self.model_id = self._detect_model_id()

    def _detect_model_id(self) -> str | None:
        """Detect model ID from command for cost calculation."""
        cli_name = self.command[0] if self.command else ""

        # Map CLI tools to model IDs for genai-prices
        # These are reasonable defaults, can be overridden in config
        if "claude" in cli_name:
            return "claude-sonnet-4-5-20250929"
        elif "gemini" in cli_name:
            return "gemini-2.0-flash-exp"
        elif "codex" in cli_name:
            return "o1"  # Codex often uses OpenAI models
        return None

    def _get_cli_type(self) -> str:
        """Detect which CLI tool this is."""
        if not self.command:
            return "unknown"
        cli_name = self.command[0].lower()
        if "claude" in cli_name:
            return "claude"
        elif "gemini" in cli_name:
            return "gemini"
        elif "codex" in cli_name:
            return "codex"
        return "unknown"

    def _should_pipe_stdin(self) -> bool:
        """Check if this CLI expects task via stdin instead of argument.

        Gemini and Codex both work with stdin piping.
        Claude takes the task as a command argument.

        Returns:
            True if task should be piped via stdin
        """
        cli_type = self._get_cli_type()
        return cli_type in ("gemini", "codex")

    def _load_schema(self, schema_name: str) -> dict | None:
        """Load a JSON schema file."""
        schema_path = Path(__file__).parent.parent / "schemas" / f"{schema_name}.json"
        if schema_path.exists():
            return json.loads(schema_path.read_text())
        return None

    def _add_structured_output_flags(
        self,
        cmd: list[str],
        schema_name: str | None = None,
    ) -> list[str]:
        """Add CLI-native structured output flags.

        Args:
            cmd: Base command
            schema_name: Name of schema file (e.g., 'execution', 'review')

        Returns:
            Enhanced command with structured output flags
        """
        cli_type = self._get_cli_type()

        if cli_type == "claude":
            # Claude supports --output-format json and --json-schema
            if "--output-format" not in cmd:
                cmd.extend(["--output-format", "json"])

            if schema_name:
                schema = self._load_schema(schema_name)
                if schema:
                    cmd.extend(["--json-schema", json.dumps(schema)])

        elif cli_type == "gemini":
            # Gemini supports --output-format json but it causes timeouts with stdin piping
            # For now, use text output and parse naturally
            # TODO: Investigate if --output-format json works differently
            pass

        elif cli_type == "codex":
            # Codex supports --json for JSONL output
            if "--json" not in cmd:
                cmd.append("--json")

            # Codex can also use --output-schema with a file
            # For now we'll rely on --json and parse the output

        return cmd

    def _parse_json_response(self, content: str) -> tuple[str, dict | None]:
        """Parse JSON response from CLI output.

        Returns:
            Tuple of (content, parsed_json)
            - content: The text content or stringified JSON
            - parsed_json: The parsed JSON object if available
        """
        cli_type = self._get_cli_type()

        try:
            if cli_type == "codex":
                # Codex outputs JSONL, collect message and usage
                lines = [line for line in content.split('\n') if line.strip()]
                message_text = None
                usage_data = None

                for line in lines:
                    try:
                        event = json.loads(line)

                        # Skip if not a dict
                        if not isinstance(event, dict):
                            continue

                        # Collect agent message
                        if event.get("type") == "item.completed":
                            item = event.get("item", {})
                            if isinstance(item, dict) and item.get("type") == "agent_message":
                                message_text = item.get("text")

                        # Collect usage info
                        if event.get("type") == "turn.completed":
                            usage_data = event.get("usage", {})
                    except json.JSONDecodeError:
                        continue

                # Combine message and usage into response
                if message_text:
                    response_data = {"message": message_text}
                    if usage_data:
                        response_data["usage"] = usage_data
                    return message_text, response_data

                return content, None

            elif cli_type == "claude":
                # Claude with --output-format json returns a result object
                parsed = json.loads(content)

                # Check for structured_output field (from --json-schema)
                if "structured_output" in parsed:
                    structured = parsed["structured_output"]
                    # Return the structured content as JSON string
                    return json.dumps(structured, indent=2), parsed

                # Fallback to other fields
                if "result" in parsed:
                    return parsed["result"], parsed
                elif "content" in parsed:
                    return parsed["content"], parsed
                return json.dumps(parsed, indent=2), parsed

            else:
                # Gemini and others output single JSON object
                parsed = json.loads(content)

                # Check if parsed is a dict before accessing fields
                if isinstance(parsed, dict):
                    if "content" in parsed:
                        return parsed["content"], parsed
                    elif "text" in parsed:
                        return parsed["text"], parsed
                    elif "message" in parsed:
                        return parsed["message"], parsed
                    return json.dumps(parsed, indent=2), parsed
                else:
                    # Parsed value is a primitive (int, str, bool, etc.)
                    # Return stringified version
                    return str(parsed), {"raw_value": parsed}

        except json.JSONDecodeError:
            # Not JSON or invalid JSON, return as-is
            return content, None

    def _extract_token_usage(self, parsed_json: dict | None) -> tuple[int, int, float | None]:
        """Extract input/output tokens and actual cost from parsed response.

        Returns:
            Tuple of (input_tokens, output_tokens, actual_cost_usd)
        """
        if not parsed_json:
            return 0, 0, None

        # Extract actual cost if available (Claude provides this!)
        actual_cost = parsed_json.get("total_cost_usd")

        # Try common token usage field names
        usage = parsed_json.get("usage", {})
        if usage:
            # Claude format: includes cache tokens
            input_tokens = (
                usage.get("input_tokens", 0) +
                usage.get("cache_creation_input_tokens", 0) +
                usage.get("cache_read_input_tokens", 0)
            )
            output_tokens = usage.get("output_tokens", usage.get("completion_tokens", 0))
            return input_tokens, output_tokens, actual_cost

        # Codex format
        if "tokens" in parsed_json:
            tokens = parsed_json["tokens"]
            return tokens.get("input", 0), tokens.get("output", 0), actual_cost

        return 0, 0, actual_cost

    def estimate_cost(self, tokens: int) -> float:
        """Estimate cost using genai-prices if available."""
        if not HAS_GENAI_PRICES or not self.model_id:
            # Fallback to rough estimate
            return tokens / 5000

        try:
            price_info = get_model_price(self.model_id)
            # Assume tokens is total (input + output), split 50/50 as rough estimate
            input_tokens = tokens // 2
            output_tokens = tokens - input_tokens
            cost = (
                input_tokens * price_info.input_price +
                output_tokens * price_info.output_price
            ) / 1_000_000
            return cost
        except Exception:
            # Fallback if model not found
            return tokens / 5000

    def estimate_cost_detailed(
        self,
        input_tokens: int,
        output_tokens: int,
    ) -> float:
        """Estimate cost with separate input/output token counts."""
        if not HAS_GENAI_PRICES or not self.model_id:
            return (input_tokens + output_tokens) / 5000

        try:
            price_info = get_model_price(self.model_id)
            cost = (
                input_tokens * price_info.input_price +
                output_tokens * price_info.output_price
            ) / 1_000_000
            return cost
        except Exception:
            return (input_tokens + output_tokens) / 5000

    async def call(
        self,
        prompt: str,
        *,
        system: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.7,
        working_dir: str | None = None,
        schema_name: str | None = None,
    ) -> AdapterResponse:
        """Make a completion call with optional structured output."""
        start = time.monotonic()
        self._call_count += 1

        cmd = list(self.command)

        # Add system prompt if supported
        if system and "claude" in self._get_cli_type():
            cmd.extend(["--system", system])

        # Add structured output flags
        if schema_name:
            cmd = self._add_structured_output_flags(cmd, schema_name)

        # Append prompt
        cmd.append(prompt)

        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=working_dir,
            env=self._get_env(),
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(),
                timeout=self.timeout_seconds,
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            raise TimeoutError(f"{self.name} timed out after {self.timeout_seconds}s")

        if proc.returncode != 0:
            error_msg = stderr.decode() if stderr else "Unknown error"
            raise RuntimeError(
                f"{self.name} failed with exit code {proc.returncode}: {error_msg}"
            )

        raw_content = stdout.decode()
        content, parsed_json = self._parse_json_response(raw_content)

        # Extract token usage if available
        input_tokens, output_tokens, actual_cost = self._extract_token_usage(parsed_json)
        total_tokens = input_tokens + output_tokens or self.estimate_tokens(prompt + content)

        # Store actual cost in raw_response for budget tracker
        if parsed_json and actual_cost is not None:
            parsed_json["_actual_cost_usd"] = actual_cost

        return AdapterResponse(
            content=content,
            token_usage=total_tokens,
            duration_seconds=time.monotonic() - start,
            raw_response=parsed_json,
        )

    async def run_agentic(
        self,
        task: str,
        *,
        working_dir: str,
        timeout_seconds: int = 1200,
        on_question: Callable[[str], str] | None = None,
        schema_name: str | None = "execution",
    ) -> AdapterResponse:
        """Run agentic task with structured output."""
        start = time.monotonic()
        self._call_count += 1

        cmd = list(self.command)

        # Add structured output flags (always call for format flags like --json)
        cmd = self._add_structured_output_flags(cmd, schema_name)

        # Determine if we should pipe task via stdin
        pipe_stdin = self._should_pipe_stdin()

        # Only append task if not piping via stdin
        if not pipe_stdin:
            cmd.append(task)

        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.PIPE if pipe_stdin else asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=working_dir,
            env=self._get_env(),
        )

        try:
            if pipe_stdin:
                # Send task via stdin
                stdout, stderr = await asyncio.wait_for(
                    proc.communicate(input=task.encode()),
                    timeout=timeout_seconds,
                )
            else:
                # Normal execution
                stdout, stderr = await asyncio.wait_for(
                    proc.communicate(),
                    timeout=timeout_seconds,
                )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            raise TimeoutError(f"{self.name} agentic task timed out after {timeout_seconds}s")

        raw_content = stdout.decode()
        content, parsed_json = self._parse_json_response(raw_content)

        # Extract token usage
        input_tokens, output_tokens, actual_cost = self._extract_token_usage(parsed_json)
        total_tokens = input_tokens + output_tokens or self.estimate_tokens(task + content)

        # Store actual cost in raw_response for budget tracker
        if parsed_json and actual_cost is not None:
            parsed_json["_actual_cost_usd"] = actual_cost

        return AdapterResponse(
            content=content,
            token_usage=total_tokens,
            duration_seconds=time.monotonic() - start,
            raw_response=parsed_json,
        )

    def _get_env(self) -> dict[str, str]:
        """Get merged environment variables."""
        merged_env = os.environ.copy()

        try:
            settings = DeliberateSettings()
            agent_env = get_agent_env_dict(self.name, settings)
            merged_env.update(agent_env)
        except Exception:
            pass

        merged_env.update(self.env)
        return merged_env

    @property
    def call_count(self) -> int:
        """Get call count."""
        return self._call_count

    def estimate_tokens(self, text: str) -> int:
        """Estimate tokens (~4 chars per token)."""
        return len(text) // 4
