"""Auto-patching for OpenAI and Anthropic SDKs.

Monkey-patches the create() methods at the class level so ALL client
instances are automatically tracked — no wrapper needed.
"""

from __future__ import annotations

import time
from typing import Any, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from . import AgentCost

from .event import CostEvent
from .tracker import _active_run

# Store originals so we can unpatch cleanly
_originals: dict[str, Any] = {}


def _patch_openai(ac: "AgentCost", agent_id: str) -> bool:
    """Patch openai.resources.chat.completions.Completions.create."""
    try:
        from openai.resources.chat.completions import Completions
    except ImportError:
        return False

    if "openai" in _originals:
        return True  # Already patched

    original_create = Completions.create
    _originals["openai"] = original_create

    def patched_create(self: Any, **kwargs: Any) -> Any:
        start = time.time()
        response = original_create(self, **kwargs)
        duration_ms = int((time.time() - start) * 1000)

        model = getattr(response, "model", kwargs.get("model", "unknown"))
        usage = getattr(response, "usage", None)

        tokens_in = getattr(usage, "prompt_tokens", 0) if usage else 0
        tokens_out = getattr(usage, "completion_tokens", 0) if usage else 0
        cost = ac.pricing.llm_cost(model, tokens_in, tokens_out)

        active = _active_run.get(None)
        effective_agent = active._agent_id if active else agent_id
        effective_run_id = active._run_id if active else None

        event = CostEvent(
            agent_id=effective_agent,
            run_id=effective_run_id,
            component_name=model,
            event_type="llm_call",
            cost_usd=cost,
            duration_ms=duration_ms,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
        )

        if active:
            active._events.append(event)

        ac._client.enqueue(event)
        return response

    Completions.create = patched_create  # type: ignore
    return True


def _patch_anthropic(ac: "AgentCost", agent_id: str) -> bool:
    """Patch anthropic.resources.messages.Messages.create."""
    try:
        from anthropic.resources.messages import Messages
    except ImportError:
        return False

    if "anthropic" in _originals:
        return True  # Already patched

    original_create = Messages.create
    _originals["anthropic"] = original_create

    def patched_create(self: Any, **kwargs: Any) -> Any:
        start = time.time()
        response = original_create(self, **kwargs)
        duration_ms = int((time.time() - start) * 1000)

        model = getattr(response, "model", kwargs.get("model", "unknown"))
        usage = getattr(response, "usage", None)

        tokens_in = getattr(usage, "input_tokens", 0) if usage else 0
        tokens_out = getattr(usage, "output_tokens", 0) if usage else 0
        cost = ac.pricing.llm_cost(model, tokens_in, tokens_out)

        active = _active_run.get(None)
        effective_agent = active._agent_id if active else agent_id
        effective_run_id = active._run_id if active else None

        event = CostEvent(
            agent_id=effective_agent,
            run_id=effective_run_id,
            component_name=model,
            event_type="llm_call",
            cost_usd=cost,
            duration_ms=duration_ms,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
        )

        if active:
            active._events.append(event)

        ac._client.enqueue(event)
        return response

    Messages.create = patched_create  # type: ignore
    return True


def patch_all(ac: "AgentCost", agent_id: str) -> list[str]:
    """Patch all detected providers. Returns list of patched provider names."""
    patched = []
    if _patch_openai(ac, agent_id):
        patched.append("openai")
    if _patch_anthropic(ac, agent_id):
        patched.append("anthropic")
    return patched


def unpatch_all() -> None:
    """Restore all original methods."""
    if "openai" in _originals:
        try:
            from openai.resources.chat.completions import Completions
            Completions.create = _originals.pop("openai")
        except ImportError:
            _originals.pop("openai", None)

    if "anthropic" in _originals:
        try:
            from anthropic.resources.messages import Messages
            Messages.create = _originals.pop("anthropic")
        except ImportError:
            _originals.pop("anthropic", None)
