"""KRIM agent adapter for Harbor (Terminal-Bench 2.0).

Installs KRIM inside the container and runs it with minimal tbench skill.

Target models: Claude Opus 4.6, OpenAI Codex 5.3

Usage:
    # Opus 4.6
    harbor run \\
        --dataset terminal-bench@2.0 \\
        --agent-import-path krim_harbor:KrimAgent \\
        --model anthropic/claude-opus-4-6 \\
        --n-concurrent 4

    # Codex 5.3
    harbor run \\
        --dataset terminal-bench@2.0 \\
        --agent-import-path krim_harbor:KrimAgent \\
        --model openai/codex-5.3 \\
        --n-concurrent 4
"""

from __future__ import annotations

import json
import logging
import os
import re
import shutil
import tempfile
from pathlib import Path

from harbor.agents.base import BaseAgent
from harbor.environments.base import BaseEnvironment
from harbor.models.agent.context import AgentContext

_REPO_ROOT = Path(__file__).resolve().parent.parent
_SKILL_PATH = _REPO_ROOT / ".krim" / "skills" / "tbench" / "SKILL.md"

# Default config tuned for Terminal-Bench
_TBENCH_CONFIG = {
    "provider": "claude",
    "max_turns": 500,
    "max_output_chars": 50000,
    "auto_commit": False,
    "ask_by_default": False,
    "max_tokens": 16384,
    "bash_timeout": 300,
    "token_limit": 180000,
    "allow_commands": [
        "ls", "cat", "head", "tail", "find", "grep", "rg", "wc",
        "git status", "git diff", "git log", "git branch",
        "python", "python3", "pip", "pip3",
        "gcc", "g++", "make", "cmake", "cargo", "rustc",
        "node", "npm", "npx",
        "apt-get", "apt", "dpkg",
        "curl", "wget",
        "pytest", "unittest",
        "cd", "mkdir", "cp", "mv", "rm", "chmod", "chown",
        "tar", "unzip", "gzip",
        "sed", "awk", "sort", "cut", "tr", "xargs",
        "docker", "systemctl", "service",
        "openssl", "ssh-keygen",
    ],
    "deny_patterns": [],
}

# GitHub repo URL for fallback installation
_KRIM_REPO = "https://github.com/batteryhob/KRIM.git"


def _extract_model_short(model_name: str) -> tuple[str, str]:
    """Extract provider/model for KRIM's --provider and --model flags.

    Supported models:
      - anthropic/claude-opus-4-6  → provider=claude, model=claude-opus-4-6
      - vertex_ai/claude-opus-4-6  → provider=vertex_ai, model=claude-opus-4-6
      - openai/codex-5.3           → provider=openai, model=codex-5.3
    """
    if "/" in model_name:
        provider_hint, model = model_name.split("/", 1)
    else:
        provider_hint, model = "", model_name

    # Check for vertex_ai first (before claude check)
    if "vertex_ai" in provider_hint.lower():
        return "vertex_ai", model
    elif "claude" in model.lower() or "anthropic" in provider_hint.lower():
        return "claude", model
    elif "codex" in model.lower() or "gpt" in model.lower() or "openai" in provider_hint.lower():
        return "openai", model
    else:
        return "claude", model


class KrimAgent(BaseAgent):
    """Harbor-compatible wrapper that runs KRIM inside the task container."""

    SUPPORTS_ATIF: bool = False

    _DEFAULT_MODEL = "claude-opus-4-6"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @staticmethod
    def name() -> str:
        return "krim"

    def version(self) -> str:
        return "0.3.0"

    async def setup(self, environment: BaseEnvironment) -> None:
        """Install KRIM inside the container."""
        self.logger.info("Setting up KRIM in container...")

        # Step 0: Ensure Python, pip, and SSL certificates are available
        self.logger.info("Checking Python/pip availability...")
        check_python = await environment.exec("which python3 || which python", timeout_sec=10)
        check_pip = await environment.exec("which pip3 || which pip", timeout_sec=10)

        # Install ca-certificates first (required for pip to work)
        self.logger.info("Installing ca-certificates...")
        await environment.exec(
            "apt-get update -qq && apt-get install -y -qq ca-certificates 2>/dev/null || "
            "apk add --no-cache ca-certificates 2>/dev/null || true",
            timeout_sec=120,
        )

        # Clean up apt/dpkg locks after ca-certificates install
        await environment.exec(
            "pkill -9 apt 2>/dev/null || true; "
            "pkill -9 dpkg 2>/dev/null || true; "
            "rm -f /var/lib/dpkg/lock* /var/lib/apt/lists/lock /var/cache/apt/archives/lock 2>/dev/null || true",
            timeout_sec=10,
        )

        if check_python.return_code != 0 or check_pip.return_code != 0:
            self.logger.info("Python/pip not found, installing...")
            # Retry Python installation up to 3 times (handles intermittent container issues)
            for attempt in range(3):
                install_result = await environment.exec(
                    "apt-get update -qq 2>/dev/null; "
                    "apt-get install -y -qq python3 python3-pip 2>/dev/null || "
                    "apk add --no-cache python3 py3-pip 2>/dev/null",
                    timeout_sec=300,
                )
                self.logger.info(f"Python install exit code: {install_result.return_code} (attempt {attempt + 1}/3)")
                if install_result.return_code == 0:
                    break
                if attempt < 2:
                    self.logger.warning(f"Python install failed, retrying in 5s...")
                    import asyncio
                    await asyncio.sleep(5)

            # Clean up any lingering apt/dpkg processes and locks
            # This prevents verifier from failing due to dpkg lock
            self.logger.info("Cleaning up apt/dpkg locks...")
            await environment.exec(
                "pkill -9 apt 2>/dev/null || true; "
                "pkill -9 dpkg 2>/dev/null || true; "
                "rm -f /var/lib/dpkg/lock* /var/lib/apt/lists/lock /var/cache/apt/archives/lock 2>/dev/null || true; "
                "dpkg --configure -a 2>/dev/null || true",
                timeout_sec=30,
            )

        # Determine pip command (pip3 preferred)
        pip_check = await environment.exec("which pip3", timeout_sec=10)
        pip_cmd = "pip3" if pip_check.return_code == 0 else "pip"

        # Step 1: Install KRIM into the container
        # Try upload from local repo first (fastest, most reliable)
        krim_pkg_dir = _REPO_ROOT / "krim"
        pyproject = _REPO_ROOT / "pyproject.toml"

        if krim_pkg_dir.is_dir() and pyproject.is_file():
            # Upload only essential files (krim/, pyproject.toml) to minimize size
            # Avoids uploading .git/, jobs/, __pycache__, etc.
            self.logger.info("Uploading KRIM source to container (minimal)...")

            with tempfile.TemporaryDirectory() as tmpdir:
                tmp_path = Path(tmpdir)
                # Copy only what pip needs
                shutil.copytree(krim_pkg_dir, tmp_path / "krim")
                shutil.copy2(pyproject, tmp_path / "pyproject.toml")
                # Copy README if exists (some builds need it)
                readme = _REPO_ROOT / "README.md"
                if readme.is_file():
                    shutil.copy2(readme, tmp_path / "README.md")

                await environment.upload_dir(str(tmp_path), "/tmp/krim-src")

            # Try with --break-system-packages first (pip 23.1+), fallback without
            result = await environment.exec(
                f"{pip_cmd} install --break-system-packages --ignore-installed /tmp/krim-src 2>&1",
                timeout_sec=120,
            )
            if result.return_code != 0 and "no such option" in (result.stdout or ""):
                self.logger.info("Retrying pip without --break-system-packages (old pip)...")
                result = await environment.exec(
                    f"{pip_cmd} install --ignore-installed /tmp/krim-src 2>&1",
                    timeout_sec=120,
                )
            self.logger.info(f"Local install exit code: {result.return_code}")
            if result.stdout:
                self.logger.info(f"pip output: {result.stdout[-500:]}")
        else:
            # Fallback: install from git
            self.logger.info("Installing KRIM from git...")
            await environment.exec(
                "apt-get update -qq && apt-get install -y -qq git 2>/dev/null || "
                "apk add --no-cache git 2>/dev/null",
                timeout_sec=60,
            )
            # Clean up apt/dpkg locks after git install
            await environment.exec(
                "pkill -9 apt 2>/dev/null || true; "
                "pkill -9 dpkg 2>/dev/null || true; "
                "rm -f /var/lib/dpkg/lock* /var/lib/apt/lists/lock /var/cache/apt/archives/lock 2>/dev/null || true",
                timeout_sec=10,
            )
            result = await environment.exec(
                f"{pip_cmd} install --break-system-packages git+{_KRIM_REPO} 2>&1",
                timeout_sec=180,
            )
            self.logger.info(f"Git install exit code: {result.return_code}")

        # Add common pip bin paths to PATH for verification
        path_export = "export PATH=$PATH:/usr/local/bin:/root/.local/bin:$(python3 -m site --user-base 2>/dev/null)/bin"

        # Verify krim is installed (try multiple methods)
        check = await environment.exec(f"{path_export} && krim --version")
        if check.return_code != 0:
            self.logger.error(
                f"KRIM installation failed. stdout: {check.stdout}, stderr: {check.stderr}"
            )
            raise RuntimeError("Failed to install KRIM in container")

        self.logger.info(f"KRIM installed: {(check.stdout or '').strip()}")

        # Upload minimal tbench skill
        if _SKILL_PATH.is_file():
            self.logger.info("Uploading tbench skill...")
            skill_content = _SKILL_PATH.read_text()
            await environment.exec(
                "mkdir -p /root/.krim/skills/tbench",
                timeout_sec=10,
            )
            # Write skill file via heredoc
            escaped_content = skill_content.replace("'", "'\\''")
            await environment.exec(
                f"cat > /root/.krim/skills/tbench/SKILL.md << 'SKILLEOF'\n{skill_content}\nSKILLEOF",
                timeout_sec=10,
            )
            self.logger.info("tbench skill uploaded")

        # Upload Google credentials for vertex_ai
        credentials_path = _REPO_ROOT / "credentials.json"
        if credentials_path.is_file():
            self.logger.info("Uploading Google credentials for vertex_ai...")
            credentials_content = credentials_path.read_text()
            await environment.exec(
                f"cat > /tmp/credentials.json << 'CREDEOF'\n{credentials_content}\nCREDEOF",
                timeout_sec=10,
            )
            self.logger.info("Google credentials uploaded")

        self.logger.info("KRIM setup complete")

    async def run(
        self,
        instruction: str,
        environment: BaseEnvironment,
        context: AgentContext,
    ) -> None:
        """Execute KRIM against the task instruction."""
        model_name = self.model_name or self._DEFAULT_MODEL
        provider, model = _extract_model_short(model_name)

        # Escape instruction for shell single quotes
        escaped_instruction = instruction.replace("'", "'\\''")

        # Build KRIM command with PATH setup for pip-installed binaries
        # Use --log-dir to save structured logs inside container
        container_log_dir = "/tmp/krim_logs"
        path_setup = "export PATH=$PATH:/usr/local/bin:/root/.local/bin:$(python3 -m site --user-base 2>/dev/null)/bin"
        # Note: verbose is now default, use --quiet to suppress turn logs
        cmd = (
            f"{path_setup} && krim "
            f"--provider {provider} "
            f"--model {model} "
            f"--max-turns {_TBENCH_CONFIG['max_turns']} "
            f"--log-dir {container_log_dir} "
            f"--skill tbench "
            f"--no-safety "
            f"'{escaped_instruction}'"
        )

        # Collect API keys and inject directly into command
        # (Daytona's env parameter doesn't work correctly with timeout wrapper)
        env_exports = []
        for key in ("ANTHROPIC_API_KEY", "OPENAI_API_KEY"):
            val = os.environ.get(key)
            if val:
                # Escape single quotes in value
                escaped_val = val.replace("'", "'\\''")
                env_exports.append(f"export {key}='{escaped_val}'")

        # For vertex_ai, set up Google credentials
        if provider == "vertex_ai":
            vertexai_location = os.environ.get("VERTEXAI_LOCATION", "us-east5")
            vertexai_project = os.environ.get("GOOGLE_CLOUD_PROJECT", "")
            env_exports.append(f"export VERTEXAI_LOCATION='{vertexai_location}'")
            env_exports.append("export GOOGLE_APPLICATION_CREDENTIALS='/tmp/credentials.json'")
            if vertexai_project:
                env_exports.append(f"export GOOGLE_CLOUD_PROJECT='{vertexai_project}'")
                env_exports.append(f"export ANTHROPIC_VERTEX_PROJECT_ID='{vertexai_project}'")

        if not env_exports:
            self.logger.warning(
                "No API keys found in host environment. "
                "Set ANTHROPIC_API_KEY or OPENAI_API_KEY."
            )

        # Prepend env exports to command
        if env_exports:
            cmd = " && ".join(env_exports) + " && " + cmd

        # Log command with masked API keys
        masked_cmd = re.sub(r"(ANTHROPIC_API_KEY|OPENAI_API_KEY)='[^']*'", r"\1='***'", cmd)
        self.logger.info(f"Running KRIM: {masked_cmd[:200]}...")

        # Run with generous timeout (most tbench tasks have 1h limit)
        # Retry on transient network/Daytona errors
        max_retries = 3
        for attempt in range(max_retries):
            try:
                result = await environment.exec(
                    cmd,
                    timeout_sec=3600,
                )
                break  # Success, exit retry loop
            except Exception as e:
                error_msg = str(e).lower()
                is_transient = any(x in error_msg for x in [
                    "unauthorized", "auth url", "network", "connection",
                    "timeout", "unavailable", "502", "503", "504"
                ])
                if is_transient and attempt < max_retries - 1:
                    self.logger.warning(f"Transient error (attempt {attempt + 1}/{max_retries}): {e}")
                    import asyncio
                    await asyncio.sleep(5 * (attempt + 1))  # Exponential backoff
                    continue
                raise  # Non-transient or final attempt, re-raise

        self.logger.info(f"KRIM finished with exit code {result.return_code}")

        # Log output for debugging (mask API keys)
        combined = (result.stdout or "") + (result.stderr or "")
        if combined:
            # Mask any API keys in output before logging
            masked = re.sub(r'(ANTHROPIC_API_KEY|OPENAI_API_KEY)=[^\s]+', r'\1=***', combined)
            masked = re.sub(r'sk-[a-zA-Z0-9_-]+', 'sk-***', masked)
            # Log last 1000 chars
            self.logger.debug(f"output (last 1000): {masked[-1000:]}")

        # Parse token usage from KRIM verbose output
        self._parse_token_usage(combined, context)

        # Download structured logs from container
        await self._download_logs(environment, container_log_dir, instruction, masked_cmd, result)

    async def _download_logs(
        self,
        environment: BaseEnvironment,
        container_log_dir: str,
        instruction: str,
        command: str,
        result,
    ) -> None:
        """Download structured logs from container to logs_dir."""
        if not self.logs_dir:
            self.logger.warning("logs_dir not set, skipping log download")
            return

        try:
            self.logs_dir.mkdir(parents=True, exist_ok=True)

            # Download logs generated by KRIM --log-dir
            for filename in ["prompt.txt", "messages.json", "meta.json"]:
                cat_result = await environment.exec(
                    f"cat {container_log_dir}/{filename} 2>/dev/null",
                    timeout_sec=10,
                )
                if cat_result.return_code == 0 and cat_result.stdout:
                    (self.logs_dir / filename).write_text(cat_result.stdout)

            # Also save instruction and stdout for debugging
            (self.logs_dir / "instruction.txt").write_text(instruction)
            (self.logs_dir / "stdout.txt").write_text(result.stdout or "")
            if result.stderr:
                (self.logs_dir / "stderr.txt").write_text(result.stderr)

            self.logger.info(f"Logs downloaded to {self.logs_dir}")

        except Exception as e:
            self.logger.warning(f"Failed to download logs: {e}")

    def _parse_token_usage(self, output: str, context: AgentContext) -> None:
        """Best-effort extraction of token counts from KRIM verbose output."""
        # KRIM prints stats like: "turns: 15 | tool calls: 42 | ~85,000 tokens"
        token_match = re.search(r"~([\d,]+)\s+tokens", output)
        if token_match:
            try:
                total = int(token_match.group(1).replace(",", ""))
                # Rough split: 70% input, 30% output (typical for tool-heavy agents)
                context.n_input_tokens = round(total * 0.7)
                context.n_output_tokens = round(total * 0.3)
            except ValueError:
                pass

        # Store raw stats in metadata
        turns_match = re.search(r"turns:\s*(\d+)", output)
        tools_match = re.search(r"tool calls:\s*(\d+)", output)
        if turns_match or tools_match:
            context.metadata = {
                "krim_turns": int(turns_match.group(1)) if turns_match else None,
                "krim_tool_calls": int(tools_match.group(1)) if tools_match else None,
            }
