from __future__ import annotations

import asyncio
import json
import logging
import time as _time
from pathlib import Path
from typing import Any, Awaitable, Callable

import anthropic

from epochler.agents import searcher
from epochler.config import get_settings
from epochler.contract.models import (
    ExperimentJournal,
    ExperimentResult,
    FrozenContract,
    HardwareProfile,
)

logger = logging.getLogger(__name__)

_SYSTEM_PROMPT = (Path(__file__).parent / "prompts" / "coder.md").read_text(encoding="utf-8")

_TOOLS = [
    {
        "name": "web_search",
        "description": "Search the web for current ML techniques, optimizers, architectures, library APIs, etc.",
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Search query"},
            },
            "required": ["query"],
        },
    },
]

_MAX_TOOL_ROUNDS = 3
_API_CALL_TIMEOUT = 135  # 15s buffer above SDK's 120s httpx timeout
_MAX_SEARCH_REQUESTS = 3

StatusCallback = Callable[[str], Awaitable[None]]


def _build_prompt(
    contract: FrozenContract,
    phase_name: str,
    current_code: str,
    journal: ExperimentJournal,
    hardware: HardwareProfile,
    input_artifacts: list[dict[str, Any]] | None = None,
    error_context: str | None = None,
) -> str:
    parts = [
        "## Frozen Contract\n",
        f"```json\n{contract.model_dump_json(indent=2)}\n```\n",
        "## Current Phase\n",
        f"{phase_name}\n",
        "## Hardware Profile\n",
        f"```json\n{hardware.model_dump_json(indent=2)}\n```\n",
        "## Current Phase Script\n",
        f"```python\n{current_code}\n```\n",
        "## Input Artifacts\n",
        f"```json\n{json.dumps(input_artifacts or [], indent=2)}\n```\n",
    ]

    if phase_name == "discovery":
        parts.append("## IMPORTANT: This is a DATA ACQUISITION script.\n")
        parts.append("Your job is to download/generate the dataset and save it to the path in the skeleton code.\n")
        parts.append("The `output` variable in the current script is the exact target path.\n\n")

    if phase_name == "analysis":
        parts.append("## IMPORTANT: This is an ANALYSIS script.\n")
        parts.append("Your job is to:\n")
        parts.append("1. Load the experiment journal from `$EPOCHLER_OUTPUT_DIR/../journal.json`\n")
        parts.append("2. Load the best model and test data from the input artifacts\n")
        parts.append("3. Generate visualizations saved as PNGs to `$EPOCHLER_OUTPUT_DIR/`:\n")
        parts.append("   - Metric progression across experiments\n")
        parts.append("   - Confusion matrix (classification) or residual plot (regression)\n")
        parts.append("   - Feature importance if applicable\n")
        parts.append("4. Use matplotlib with Agg backend: `import matplotlib; matplotlib.use('Agg')`\n")
        parts.append("5. Print a single JSON line to stdout with keys: `metrics` (dict), `summary` (string), `findings` (list of strings), `artifacts` (list of file paths)\n\n")

    if error_context:
        parts.append("## Error From Previous Attempt\n")
        parts.append("The previous run of this script FAILED with the following error output. ")
        parts.append("You MUST fix the issue described below. Do NOT repeat the same mistake.\n\n")
        parts.append(f"```\n{error_context}\n```\n")

    parts.append("## Experiment Journal\n")

    if journal.experiments:
        for exp in journal.experiments[-10:]:
            parts.append(f"### Run #{exp.run_number} — {exp.status.value}\n")
            parts.append(f"- Hypothesis: {exp.hypothesis}\n")
            parts.append(f"- Metrics: {json.dumps(exp.metrics)}\n")
            parts.append(f"- Verdict: {exp.verdict_reason}\n")
            parts.append(f"- Lessons: {exp.lessons_learned}\n")
            parts.append(f"- Next: {exp.next_direction}\n\n")
    else:
        parts.append("No experiments yet. This is the first run.\n")

    from datetime import date as _date
    parts.append(f"\n## Today's Date\n{_date.today().strftime('%A, %B %d, %Y')}\n")

    parts.append(
        "\nNow produce your next iteration. Return valid JSON with keys: "
        "code, hypothesis, description, search_requests (optional list).\n"
    )
    return "\n".join(parts)


async def _handle_tool_call(tool_name: str, tool_input: dict[str, Any]) -> str:
    if tool_name == "web_search":
        query = tool_input.get("query", "")
        if not query:
            return json.dumps({"error": "Empty search query"})
        result = await searcher.search(query)
        return json.dumps(result, indent=2)
    return f"Unknown tool: {tool_name}"


async def generate_iteration(
    contract: FrozenContract,
    current_code: str,
    journal: ExperimentJournal,
    hardware: HardwareProfile,
    input_artifacts: list[dict[str, Any]] | None = None,
    *,
    model_override: str | None = None,
    error_context: str | None = None,
    on_status: StatusCallback | None = None,
    watchdog: Any | None = None,
) -> dict[str, Any]:
    """Ask the coder agent to produce the next train.py iteration."""
    return await generate_phase_script(
        contract,
        phase_name="training",
        current_code=current_code,
        journal=journal,
        hardware=hardware,
        input_artifacts=input_artifacts,
        model_override=model_override,
        error_context=error_context,
        on_status=on_status,
        watchdog=watchdog,
    )


async def generate_phase_script(
    contract: FrozenContract,
    phase_name: str,
    current_code: str,
    journal: ExperimentJournal,
    hardware: HardwareProfile,
    input_artifacts: list[dict[str, Any]] | None = None,
    *,
    model_override: str | None = None,
    error_context: str | None = None,
    on_status: StatusCallback | None = None,
    watchdog: Any | None = None,
) -> dict[str, Any]:
    """Ask the coder agent to produce the next phase script."""
    settings = get_settings()
    model = model_override or settings.anthropic_model
    text_parts: list[str] = []

    logger.info(
        "Coder: generate_phase_script START phase=%s model=%s error_context=%s",
        phase_name, model, "yes" if error_context else "no",
    )

    async with anthropic.AsyncAnthropic(
        api_key=settings.anthropic_api_key,
        timeout=anthropic.Timeout(timeout=120.0, connect=10.0),
        max_retries=0,
    ) as client:
        logger.debug("Coder: AsyncAnthropic client opened")
        if watchdog:
            watchdog.register_client(client)
        try:
            user_prompt = _build_prompt(
                contract, phase_name, current_code, journal, hardware, input_artifacts,
                error_context=error_context,
            )
            messages: list[dict[str, Any]] = [{"role": "user", "content": user_prompt}]
            logger.debug("Coder: prompt built, %d chars", len(user_prompt))
            tool_round = 0

            while tool_round <= _MAX_TOOL_ROUNDS:
                if on_status:
                    try:
                        await on_status(f"Calling coder API (round {tool_round + 1})...")
                    except Exception:
                        pass

                if watchdog:
                    watchdog.heartbeat(f"api_call:coder:{phase_name}:round_{tool_round}")

                t0 = _time.monotonic()
                logger.info("Coder: API call START round=%d model=%s", tool_round + 1, model)
                try:
                    response = await asyncio.wait_for(
                        client.messages.create(
                            model=model,
                            max_tokens=8192,
                            system=_SYSTEM_PROMPT,
                            messages=messages,
                            tools=_TOOLS,
                        ),
                        timeout=_API_CALL_TIMEOUT,
                    )
                except asyncio.TimeoutError:
                    logger.warning("Coder API call TIMED OUT (round %d, %.1fs)", tool_round + 1, _time.monotonic() - t0)
                    break
                except asyncio.CancelledError:
                    logger.warning("Coder API call CANCELLED (round %d, %.1fs)", tool_round + 1, _time.monotonic() - t0)
                    raise
                except Exception:
                    logger.exception("Coder API call FAILED (round %d, %.1fs)", tool_round + 1, _time.monotonic() - t0)
                    break

                elapsed = _time.monotonic() - t0
                logger.info(
                    "Coder: API call DONE round=%d elapsed=%.1fs stop=%s tokens_in=%s tokens_out=%s",
                    tool_round + 1, elapsed,
                    getattr(response, "stop_reason", "?"),
                    getattr(getattr(response, "usage", None), "input_tokens", "?"),
                    getattr(getattr(response, "usage", None), "output_tokens", "?"),
                )
                if watchdog:
                    watchdog.heartbeat(f"processing:coder:{phase_name}:round_{tool_round}")

                tool_uses = []
                for block in response.content:
                    if block.type == "text":
                        text_parts.append(block.text)
                    elif block.type == "tool_use":
                        tool_uses.append(block)

                if not tool_uses:
                    break

                tool_round += 1
                if tool_round > _MAX_TOOL_ROUNDS:
                    logger.warning("Coder hit tool-use cap (%d rounds), using collected text", _MAX_TOOL_ROUNDS)
                    break

                messages.append({"role": "assistant", "content": response.content})
                tool_results = []
                for i, tool_use in enumerate(tool_uses):
                    if on_status:
                        try:
                            await on_status(f"Tool call {i + 1}/{len(tool_uses)}: {tool_use.name}...")
                        except Exception:
                            pass
                    if watchdog:
                        watchdog.heartbeat(f"tool:{tool_use.name}:round_{tool_round}")
                    try:
                        result_text = await _handle_tool_call(tool_use.name, tool_use.input)
                    except Exception as exc:
                        result_text = json.dumps({"error": str(exc)})
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": result_text,
                    })
                messages.append({"role": "user", "content": tool_results})
        finally:
            if watchdog:
                watchdog.clear_client()
            logger.debug("Coder: AsyncAnthropic client closed")

    raw = "\n".join(text_parts).strip()
    logger.info("Coder: collected %d text parts, raw length=%d", len(text_parts), len(raw))

    try:
        if "```json" in raw:
            raw = raw.split("```json", 1)[1].split("```", 1)[0].strip()
        elif "```" in raw:
            raw = raw.split("```", 1)[1].split("```", 1)[0].strip()
        result = json.loads(raw)
    except (json.JSONDecodeError, IndexError):
        logger.error("Failed to parse coder response as JSON: %s", raw[:500])
        result = {
            "code": current_code,
            "hypothesis": "Parse error — no changes made",
            "description": "Failed to parse agent output",
            "search_requests": [],
        }

    raw_requests = result.get("search_requests", [])
    if isinstance(raw_requests, str):
        raw_requests = [raw_requests]
    if not isinstance(raw_requests, list):
        raw_requests = []
    raw_requests = [q for q in raw_requests if isinstance(q, str) and q.strip()][:_MAX_SEARCH_REQUESTS]

    if watchdog:
        watchdog.heartbeat(f"searches:coder:{phase_name}")

    search_results = []
    for query in raw_requests:
        sr = await searcher.search(query)
        search_results.append({"query": query, "summary": sr.get("summary", "")})
    result["search_queries"] = search_results

    return result
