"""
Patcher for Pydantic AI.
Wraps: Agent.run, Agent.run_sync
"""

import time
import functools
import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ..emitter import NotiLensEmitter
    from ..loop_detector import LoopDetector

logger = logging.getLogger("notilens.pydantic_ai")

_patched = False


def patch(emitter: "NotiLensEmitter", loop_detector: "LoopDetector", timeout: float) -> bool:
    global _patched
    if _patched:
        return True
    try:
        from pydantic_ai import Agent

        # ── run_sync ──────────────────────────────────────────────────────
        _orig_sync = Agent.run_sync

        @functools.wraps(_orig_sync)
        def _run_sync(self, user_prompt, *args, **kwargs):
            model = str(getattr(self, "model", "unknown"))
            start = time.time()

            loop_msg = loop_detector.record()
            if loop_msg:
                emitter.send("task.loop", loop_msg, level="warning", meta={"model": model, "source": "pydantic-ai"})

            emitter.send("ai.call.start", f"PydanticAI run — {model}",
                         level="info", meta={"model": model, "prompt_len": len(str(user_prompt)), "source": "pydantic-ai"})
            try:
                result = _orig_sync(self, user_prompt, *args, **kwargs)
                duration = round(time.time() - start, 3)
                if duration > timeout:
                    emitter.send("task.timeout", f"PydanticAI run took {duration}s",
                                 level="warning", meta={"duration": duration, "source": "pydantic-ai"})
                emitter.send("ai.call.complete", f"PydanticAI done in {duration}s",
                             level="info", meta={"model": model, "duration": duration, "source": "pydantic-ai"})
                return result
            except Exception as exc:
                emitter.send("task.error", f"PydanticAI error: {exc}",
                             level="error", meta={"model": model, "error": str(exc), "source": "pydantic-ai"})
                raise

        Agent.run_sync = _run_sync

        # ── async run ─────────────────────────────────────────────────────
        _orig_run = Agent.run

        @functools.wraps(_orig_run)
        async def _run_async(self, user_prompt, *args, **kwargs):
            model = str(getattr(self, "model", "unknown"))
            start = time.time()

            loop_msg = loop_detector.record()
            if loop_msg:
                emitter.send("task.loop", loop_msg, level="warning", meta={"model": model, "source": "pydantic-ai"})

            emitter.send("ai.call.start", f"PydanticAI async — {model}",
                         level="info", meta={"model": model, "source": "pydantic-ai"})
            try:
                result = await _orig_run(self, user_prompt, *args, **kwargs)
                duration = round(time.time() - start, 3)
                emitter.send("ai.call.complete", f"PydanticAI async done in {duration}s",
                             level="info", meta={"model": model, "duration": duration, "source": "pydantic-ai"})
                return result
            except Exception as exc:
                emitter.send("task.error", f"PydanticAI async error: {exc}",
                             level="error", meta={"model": model, "error": str(exc), "source": "pydantic-ai"})
                raise

        Agent.run = _run_async

        _patched = True
        logger.debug("NotiLens: Pydantic AI patched")
        return True

    except ImportError:
        return False
    except Exception as e:
        logger.warning("NotiLens: failed to patch Pydantic AI — %s", e)
        return False
