import asyncio
import time
from datetime import UTC, datetime
from typing import Any

import boto3
import structlog

from marketing_intelligence.agent.mcp_client import MCPToolClient
from marketing_intelligence.core.config import settings

logger = structlog.get_logger("agent.bedrock")


def _key_inputs(name: str, inputs: dict[str, Any]) -> dict[str, Any]:
    if name in ("scrape_sound_page", "build_sound_url"):
        return {"sound_id": inputs.get("sound_id")}
    if name in ("scrape_video", "scrape_url"):
        return {"url": inputs.get("url")}
    if name == "scrape_tag_page":
        return {"tag": inputs.get("tag")}
    if name == "export_results":
        return {"label": inputs.get("label")}
    if name in ("close_browser_session", "apply_behavior_tactics"):
        return {"session_id": inputs.get("session_id")}
    return {}


def _key_result(name: str, result: Any) -> dict[str, Any]:
    if not isinstance(result, dict):
        return {}
    ok = result.get("success", True)
    base = {"ok": ok}
    if not ok:
        return {**base, "error": str(result.get("error", ""))[:120]}
    if name == "scrape_tag_page":
        return {
            **base,
            "videos": len(result.get("videos", [])),
            "video_count_text": result.get("video_count_text"),
        }
    if name == "scrape_sound_page":
        return {
            **base,
            "title": result.get("title"),
            "video_urls": len(result.get("video_urls", [])),
            "video_count_text": result.get("video_count_text"),
        }
    if name == "scrape_video":
        return {
            **base,
            "views": result.get("views"),
            "likes": result.get("likes"),
            "sound_id": result.get("sound_id"),
        }
    if name == "scrape_url":
        return {**base, "chars": len(str(result.get("content", "")))}
    if name == "export_results":
        return {**base, "path": result.get("saved_to"), "s3_uri": result.get("s3_uri")}
    block = (result.get("signals") or {}).get("block_signal")
    if block and block not in ("none", None):
        base["block_signal"] = block
    return base


def _prune_old_results(messages: list[dict[str, Any]], keep_last: int = 6) -> None:
    result_items = [
        item
        for msg in messages
        if msg["role"] == "user" and isinstance(msg["content"], list)
        for item in msg["content"]
        if isinstance(item, dict) and "toolResult" in item
    ]
    for item in result_items[:-keep_last]:
        item["toolResult"]["content"] = [{"text": "ok"}]


class BedrockAgent:
    def __init__(self, model_id: str | None = None) -> None:
        self._client = boto3.client("bedrock-runtime", region_name=settings.aws_region)
        self._model_id = model_id or settings.bedrock_model_id

    async def run(
        self,
        task: str,
        params: dict[str, Any] | None = None,
        run_id: str | None = None,
        allowed_tools: set[str] | None = None,
    ) -> str:
        messages: list[dict[str, Any]] = [{"role": "user", "content": [{"text": task}]}]
        step = 0
        agent_start = time.time()
        run_id = run_id or datetime.now(UTC).strftime("%Y%m%d_%H%M%S")

        log = logger.bind(model=self._model_id, backend="bedrock")
        log.info("agent.start", task=task[:200], params=params or {})

        async with MCPToolClient() as mcp:
            tool_specs = await mcp.list_tools_bedrock(allowed=allowed_tools)
            system_prompt = await mcp.get_prompt("tiktok_system", params or {})

            while True:
                _prune_old_results(messages)
                response = await asyncio.to_thread(
                    self._client.converse,
                    modelId=self._model_id,
                    system=[
                        {"text": system_prompt},
                        {"cachePoint": {"type": "default"}},
                    ],
                    messages=messages,
                    toolConfig={"tools": tool_specs},
                )

                output_message = response["output"]["message"]
                messages.append(output_message)
                stop_reason = response["stopReason"]

                if stop_reason == "end_turn":
                    elapsed_ms = int((time.time() - agent_start) * 1000)
                    log.info("agent.done", steps=step, elapsed_ms=elapsed_ms)
                    text_blocks = [
                        b["text"] for b in output_message["content"] if "text" in b
                    ]
                    return text_blocks[0] if text_blocks else ""

                if stop_reason != "tool_use":
                    return f"Unexpected stop reason: {stop_reason}"

                # Converse API content blocks are single-key dicts: {"text": ...} or
                # {"toolUse": {"toolUseId", "name", "input"}} — there is no "type" field.
                tool_results = []
                for block in output_message["content"]:
                    tool_use = block.get("toolUse")
                    if not tool_use:
                        continue
                    step += 1
                    name = tool_use["name"]
                    key_in = _key_inputs(name, tool_use["input"] or {})
                    log.info("tool.call", step=step, tool=name, **key_in)

                    t0 = time.time()
                    try:
                        result = await mcp.call_tool(
                            name, dict(tool_use["input"] or {})
                        )
                    except Exception as exc:
                        result = {"error": str(exc), "success": False}

                    elapsed_ms = int((time.time() - t0) * 1000)
                    key_out = _key_result(name, result)
                    log.info(
                        "tool.result",
                        step=step,
                        tool=name,
                        elapsed_ms=elapsed_ms,
                        **key_out,
                    )

                    is_error = (
                        isinstance(result, dict) and result.get("success") is False
                    )
                    tool_results.append(
                        {
                            "toolResult": {
                                "toolUseId": tool_use["toolUseId"],
                                "content": [
                                    {
                                        "json": result
                                        if isinstance(result, dict)
                                        else {"value": result}
                                    }
                                ],
                                "status": "error" if is_error else "success",
                            }
                        }
                    )

                if not tool_results:
                    log.error("agent.no_tool_use_blocks", stop_reason=stop_reason)
                    return "Bedrock returned stop_reason=tool_use but no toolUse blocks were found"

                messages.append({"role": "user", "content": tool_results})