"""Auto-instrumentation for OpenAI and Anthropic clients."""
from __future__ import annotations

import functools
import time
import uuid
from typing import Optional

from art._config import get_config
from art._context import (
    add_budget_spent,
    get_budget_limit,
    get_budget_spent,
    get_step_index,
    get_trace_id,
    increment_step,
)
from art._cost import BudgetExceededError, estimate_cost, estimate_tokens, _PRICES_PER_1K_INPUT, _DEFAULT_PRICE_PER_1K
from art._evaluators import evaluate
from art._exporters import make_exporter
from art._record import EvaluationRecord


def _build_record(
    *,
    model: str,
    pipeline: str,
    prompt: str,
    response: str,
    latency_ms: int,
    prompt_tokens: int,
    completion_tokens: int,
    estimated_cost: float,
    eval_result,
    trace_id: Optional[str],
    step_index: int,
) -> EvaluationRecord:
    return EvaluationRecord(
        id=str(uuid.uuid4()),
        created_at=EvaluationRecord.now_iso(),
        model=model,
        pipeline=pipeline,
        prompt=prompt,
        response=response,
        latency_ms=latency_ms,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        estimated_cost_usd=estimated_cost,
        confidence_score=eval_result.confidence_score,
        hallucination_score=eval_result.hallucination_score,
        coherence_score=eval_result.coherence_score,
        instruction_following_score=eval_result.instruction_following_score,
        hallucination_indicators=eval_result.hallucination_indicators,
        all_indicators=eval_result.all_indicators,
        trace_id=trace_id,
        step_index=step_index,
        source="sdk",
    )


def _extract_prompt(messages: list[dict]) -> str:
    parts = []
    for m in messages:
        content = m.get("content", "")
        if isinstance(content, str):
            parts.append(content)
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    parts.append(block.get("text", ""))
    return "\n".join(parts)


def instrument_openai(client=None) -> None:
    """Patch the OpenAI chat completions API to automatically evaluate every call.

    Call once at startup::

        import art, openai
        art.instrument_openai()
        client = openai.OpenAI()

    Or pass a specific client instance to patch only that instance::

        art.instrument_openai(client=my_client)
    """
    try:
        import openai as _openai
        from openai.resources.chat.completions import Completions
    except ImportError:
        raise ImportError(
            "openai package is not installed. "
            "Install it with: pip install art-sdk[openai]"
        )

    if client is not None:
        # Patch a specific client instance
        original = client.chat.completions.create
        client.chat.completions.create = _make_sync_wrapper(original)
        return

    # Patch globally at the class level
    if getattr(Completions, "_art_patched", False):
        return  # Already patched — idempotent

    original_create = Completions.create

    @functools.wraps(original_create)
    def patched_create(self_inner, **kwargs):
        return _sync_wrapper(original_create, self_inner, **kwargs)

    Completions.create = patched_create
    Completions._art_patched = True  # type: ignore[attr-defined]

    # Also patch async
    try:
        from openai.resources.chat.completions import AsyncCompletions

        if not getattr(AsyncCompletions, "_art_patched", False):
            original_acreate = AsyncCompletions.create

            @functools.wraps(original_acreate)
            async def patched_acreate(self_inner, **kwargs):
                return await _async_wrapper(original_acreate, self_inner, **kwargs)

            AsyncCompletions.create = patched_acreate
            AsyncCompletions._art_patched = True  # type: ignore[attr-defined]
    except (ImportError, AttributeError):
        pass


def _make_sync_wrapper(original):
    @functools.wraps(original)
    def wrapper(**kwargs):
        return _sync_wrapper(original, None, **kwargs)
    return wrapper


def _sync_wrapper(original, self_inner, **kwargs):
    config = get_config()
    messages = kwargs.get("messages", [])
    model = kwargs.get("model", "")

    # Pre-call budget check
    estimated_cost = estimate_cost(messages, model)
    budget_limit = get_budget_limit()
    if budget_limit is not None:
        spent = get_budget_spent()
        if spent + estimated_cost > budget_limit:
            raise BudgetExceededError(
                f"Request would cost ~${estimated_cost:.5f}, but only "
                f"${budget_limit - spent:.5f} remains in budget "
                f"(limit: ${budget_limit:.4f})",
                estimated_tokens=estimate_tokens(messages, model),
                estimated_usd=estimated_cost,
                limit_usd=budget_limit,
            )

    # Global per-request guard
    if config.max_usd_per_request and estimated_cost > config.max_usd_per_request:
        raise BudgetExceededError(
            f"Request estimated cost ${estimated_cost:.5f} exceeds "
            f"per-request limit ${config.max_usd_per_request:.5f}",
            estimated_tokens=estimate_tokens(messages, model),
            estimated_usd=estimated_cost,
            limit_usd=config.max_usd_per_request,
        )

    trace_id = get_trace_id()
    step_index = increment_step() if trace_id else 0

    start = time.monotonic()
    if self_inner is not None:
        result = original(self_inner, **kwargs)
    else:
        result = original(**kwargs)
    latency_ms = int((time.monotonic() - start) * 1000)

    try:
        _process_result(
            result=result,
            messages=messages,
            model=model,
            config=config,
            estimated_cost=estimated_cost,
            latency_ms=latency_ms,
            trace_id=trace_id,
            step_index=step_index,
        )
    except Exception:
        pass  # Evaluation errors must never affect the calling application

    if budget_limit is not None:
        add_budget_spent(estimated_cost)

    return result


async def _async_wrapper(original, self_inner, **kwargs):
    config = get_config()
    messages = kwargs.get("messages", [])
    model = kwargs.get("model", "")

    estimated_cost = estimate_cost(messages, model)
    budget_limit = get_budget_limit()
    if budget_limit is not None:
        spent = get_budget_spent()
        if spent + estimated_cost > budget_limit:
            raise BudgetExceededError(
                f"Request would cost ~${estimated_cost:.5f}, but only "
                f"${budget_limit - spent:.5f} remains in budget",
                estimated_tokens=estimate_tokens(messages, model),
                estimated_usd=estimated_cost,
                limit_usd=budget_limit,
            )

    if config.max_usd_per_request and estimated_cost > config.max_usd_per_request:
        raise BudgetExceededError(
            f"Request estimated cost ${estimated_cost:.5f} exceeds "
            f"per-request limit ${config.max_usd_per_request:.5f}",
            estimated_tokens=estimate_tokens(messages, model),
            estimated_usd=estimated_cost,
            limit_usd=config.max_usd_per_request,
        )

    trace_id = get_trace_id()
    step_index = increment_step() if trace_id else 0

    start = time.monotonic()
    result = await original(self_inner, **kwargs)
    latency_ms = int((time.monotonic() - start) * 1000)

    try:
        _process_result(
            result=result,
            messages=messages,
            model=model,
            config=config,
            estimated_cost=estimated_cost,
            latency_ms=latency_ms,
            trace_id=trace_id,
            step_index=step_index,
        )
    except Exception:
        pass

    if budget_limit is not None:
        add_budget_spent(estimated_cost)

    return result


def _process_result(*, result, messages, model, config, estimated_cost, latency_ms, trace_id, step_index):
    choices = getattr(result, "choices", []) or []
    response_text = ""
    if choices:
        msg = getattr(choices[0], "message", None)
        if msg:
            response_text = getattr(msg, "content", "") or ""

    usage = getattr(result, "usage", None)
    prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
    completion_tokens = getattr(usage, "completion_tokens", 0) or 0

    # Use real token count for cost if available
    if prompt_tokens:
        price = _PRICES_PER_1K_INPUT.get(model, _DEFAULT_PRICE_PER_1K)
        actual_cost = round((prompt_tokens / 1000) * price, 8)
    else:
        actual_cost = estimated_cost

    prompt = _extract_prompt(messages)
    eval_result = evaluate(prompt, response_text)

    record = _build_record(
        model=model,
        pipeline=config.pipeline,
        prompt=prompt,
        response=response_text,
        latency_ms=latency_ms,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        estimated_cost=actual_cost,
        eval_result=eval_result,
        trace_id=trace_id,
        step_index=step_index,
    )

    exporter = make_exporter(config)
    exporter.export(record)


# ── Anthropic instrumentation ─────────────────────────────────────────────────

def instrument_anthropic(client=None) -> None:
    """Patch the Anthropic messages API to automatically evaluate every call.

    Call once at startup::

        import art, anthropic
        art.instrument_anthropic()
        client = anthropic.Anthropic()

    Or pass a specific client instance::

        art.instrument_anthropic(client=my_client)
    """
    try:
        from anthropic.resources.messages import Messages, AsyncMessages
    except ImportError:
        raise ImportError(
            "anthropic package is not installed. "
            "Install it with: pip install art-sdk[anthropic]"
        )

    if client is not None:
        original = client.messages.create
        client.messages.create = _make_anthropic_sync_wrapper(original)
        return

    if getattr(Messages, "_art_patched", False):
        return

    original_create = Messages.create

    @functools.wraps(original_create)
    def patched_create(self_inner, **kwargs):
        return _anthropic_sync_wrapper(original_create, self_inner, **kwargs)

    Messages.create = patched_create
    Messages._art_patched = True  # type: ignore[attr-defined]

    try:
        if not getattr(AsyncMessages, "_art_patched", False):
            original_acreate = AsyncMessages.create

            @functools.wraps(original_acreate)
            async def patched_acreate(self_inner, **kwargs):
                return await _anthropic_async_wrapper(original_acreate, self_inner, **kwargs)

            AsyncMessages.create = patched_acreate
            AsyncMessages._art_patched = True  # type: ignore[attr-defined]
    except (ImportError, AttributeError):
        pass


def _make_anthropic_sync_wrapper(original):
    @functools.wraps(original)
    def wrapper(**kwargs):
        return _anthropic_sync_wrapper(original, None, **kwargs)
    return wrapper


def _anthropic_extract_prompt(kwargs: dict) -> str:
    """Combine system prompt + user messages into a single string for evaluation."""
    parts = []
    system = kwargs.get("system", "")
    if system:
        parts.append(system)
    for m in kwargs.get("messages", []):
        content = m.get("content", "")
        if isinstance(content, str):
            parts.append(content)
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    parts.append(block.get("text", ""))
    return "\n".join(parts)


def _anthropic_extract_response(result) -> str:
    """Extract text from an Anthropic Message response."""
    content = getattr(result, "content", []) or []
    for block in content:
        if getattr(block, "type", None) == "text":
            return getattr(block, "text", "") or ""
    return ""


def _anthropic_sync_wrapper(original, self_inner, **kwargs):
    config = get_config()
    model = kwargs.get("model", "")
    messages = kwargs.get("messages", [])

    estimated_cost = estimate_cost(messages, model)
    budget_limit = get_budget_limit()
    if budget_limit is not None:
        spent = get_budget_spent()
        if spent + estimated_cost > budget_limit:
            raise BudgetExceededError(
                f"Request would cost ~${estimated_cost:.5f}, but only "
                f"${budget_limit - spent:.5f} remains in budget (limit: ${budget_limit:.4f})",
                estimated_tokens=estimate_tokens(messages, model),
                estimated_usd=estimated_cost,
                limit_usd=budget_limit,
            )

    if config.max_usd_per_request and estimated_cost > config.max_usd_per_request:
        raise BudgetExceededError(
            f"Request estimated cost ${estimated_cost:.5f} exceeds "
            f"per-request limit ${config.max_usd_per_request:.5f}",
            estimated_tokens=estimate_tokens(messages, model),
            estimated_usd=estimated_cost,
            limit_usd=config.max_usd_per_request,
        )

    trace_id = get_trace_id()
    step_index = increment_step() if trace_id else 0

    start = time.monotonic()
    if self_inner is not None:
        result = original(self_inner, **kwargs)
    else:
        result = original(**kwargs)
    latency_ms = int((time.monotonic() - start) * 1000)

    try:
        _process_anthropic_result(
            result=result,
            kwargs=kwargs,
            model=model,
            config=config,
            estimated_cost=estimated_cost,
            latency_ms=latency_ms,
            trace_id=trace_id,
            step_index=step_index,
        )
    except Exception:
        pass

    if budget_limit is not None:
        add_budget_spent(estimated_cost)

    return result


async def _anthropic_async_wrapper(original, self_inner, **kwargs):
    config = get_config()
    model = kwargs.get("model", "")
    messages = kwargs.get("messages", [])

    estimated_cost = estimate_cost(messages, model)
    budget_limit = get_budget_limit()
    if budget_limit is not None:
        spent = get_budget_spent()
        if spent + estimated_cost > budget_limit:
            raise BudgetExceededError(
                f"Request would cost ~${estimated_cost:.5f}, but only "
                f"${budget_limit - spent:.5f} remains in budget",
                estimated_tokens=estimate_tokens(messages, model),
                estimated_usd=estimated_cost,
                limit_usd=budget_limit,
            )

    if config.max_usd_per_request and estimated_cost > config.max_usd_per_request:
        raise BudgetExceededError(
            f"Request estimated cost ${estimated_cost:.5f} exceeds "
            f"per-request limit ${config.max_usd_per_request:.5f}",
            estimated_tokens=estimate_tokens(messages, model),
            estimated_usd=estimated_cost,
            limit_usd=config.max_usd_per_request,
        )

    trace_id = get_trace_id()
    step_index = increment_step() if trace_id else 0

    start = time.monotonic()
    result = await original(self_inner, **kwargs)
    latency_ms = int((time.monotonic() - start) * 1000)

    try:
        _process_anthropic_result(
            result=result,
            kwargs=kwargs,
            model=model,
            config=config,
            estimated_cost=estimated_cost,
            latency_ms=latency_ms,
            trace_id=trace_id,
            step_index=step_index,
        )
    except Exception:
        pass

    if budget_limit is not None:
        add_budget_spent(estimated_cost)

    return result


def _process_anthropic_result(*, result, kwargs, model, config, estimated_cost, latency_ms, trace_id, step_index):
    response_text = _anthropic_extract_response(result)
    prompt_text = _anthropic_extract_prompt(kwargs)

    usage = getattr(result, "usage", None)
    input_tokens = getattr(usage, "input_tokens", 0) or 0
    output_tokens = getattr(usage, "output_tokens", 0) or 0

    if input_tokens:
        price = _PRICES_PER_1K_INPUT.get(model, _DEFAULT_PRICE_PER_1K)
        actual_cost = round((input_tokens / 1000) * price, 8)
    else:
        actual_cost = estimated_cost

    eval_result = evaluate(prompt_text, response_text)

    record = _build_record(
        model=model,
        pipeline=config.pipeline,
        prompt=prompt_text,
        response=response_text,
        latency_ms=latency_ms,
        prompt_tokens=input_tokens,
        completion_tokens=output_tokens,
        estimated_cost=actual_cost,
        eval_result=eval_result,
        trace_id=trace_id,
        step_index=step_index,
    )

    exporter = make_exporter(config)
    exporter.export(record)
