from __future__ import annotations

import logging
from pathlib import Path
from uuid import uuid4

import typer
from pydantic import ValidationError

from dr_llm.logging import emit_generation_event, generation_log_context
from dr_llm.providers import build_default_registry
from dr_llm.providers.llm_request import LlmRequest
from dr_llm.providers.reasoning import ReasoningConfig
from dr_llm.pool.db import PoolDb

from . import common


def register(app: typer.Typer) -> None:
    app.command("query")(query)


def query(
    provider: str = typer.Option(..., help="Provider key registered in dr-llm."),
    model: str = typer.Option(..., help="Model identifier for the provider."),
    message: list[str] = typer.Option(
        None, "--message", help="User message. Repeatable."
    ),
    messages_file: Path | None = typer.Option(
        None, help="Path to JSON messages payload."
    ),
    temperature: float | None = typer.Option(None),
    top_p: float | None = typer.Option(None),
    max_tokens: int | None = typer.Option(None),
    reasoning_json: str | None = typer.Option(
        None,
        help='JSON reasoning config (e.g. {"effort":"high"} or {"max_tokens":2000}).',
    ),
    metadata_json: str | None = typer.Option(None, help="JSON object metadata."),
    run_id: str | None = typer.Option(None),
    external_call_id: str | None = typer.Option(None),
    record: bool = typer.Option(True, "--record/--no-record"),
    dsn: str | None = typer.Option(None, envvar="DR_LLM_DATABASE_URL"),
    min_pool_size: int = typer.Option(4),
    max_pool_size: int = typer.Option(64),
) -> None:
    """Execute a single LLM query through the unified provider interface."""
    metadata = (
        common._parse_json(metadata_json, arg_name="metadata_json", expected=dict) or {}
    )
    reasoning_payload = common._parse_json(
        reasoning_json, arg_name="reasoning_json", expected=dict
    )
    try:
        reasoning = (
            ReasoningConfig(**reasoning_payload)
            if isinstance(reasoning_payload, dict)
            else None
        )
    except ValidationError as exc:
        raise typer.BadParameter(str(exc)) from exc
    messages_payload = common._load_messages(messages_file, message or [])

    try:
        request = LlmRequest(
            provider=provider,
            model=model,
            messages=messages_payload,
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            reasoning=reasoning,
            metadata=metadata,
        )
    except ValidationError as exc:
        raise typer.BadParameter(str(exc)) from exc

    registry = build_default_registry()
    try:
        adapter = registry.get(provider)
        call_id = external_call_id or uuid4().hex

        repository: PoolDb | None = None
        if record:
            repository = common._repo(dsn, min_pool_size, max_pool_size)

        log_context = {
            "call_id": call_id,
            "run_id": run_id,
            "provider": request.provider,
            "model": request.model,
            "mode": adapter.mode,
        }
        with generation_log_context(log_context):
            emit_generation_event(
                event_type="llm_call.started",
                stage="query.before_adapter",
                payload={
                    "request": request.model_dump(
                        mode="json",
                        exclude_none=True,
                        exclude_computed_fields=True,
                    )
                },
            )
            try:
                response = adapter.generate(request)
            except Exception as exc:  # noqa: BLE001
                emit_generation_event(
                    event_type="llm_call.failed",
                    stage="query.adapter_exception",
                    payload={
                        "error_type": type(exc).__name__,
                        "message": str(exc),
                    },
                )
                if repository is not None:
                    try:
                        repository.record_call(
                            request=request,
                            response=None,
                            run_id=run_id,
                            status="failed",
                            mode=adapter.mode,
                            error_text=str(exc),
                            external_call_id=external_call_id,
                            metadata=metadata,
                            call_id=call_id,
                        )
                    except Exception as rec_exc:  # noqa: BLE001
                        logging.getLogger(__name__).warning(
                            "Failed to record failed call (call_id=%s, run_id=%s): %s",
                            call_id,
                            run_id,
                            rec_exc,
                        )
                raise

            emit_generation_event(
                event_type="llm_call.succeeded",
                stage="query.after_adapter",
                payload={
                    "response": response.model_dump(
                        mode="json",
                        exclude_none=True,
                        exclude_computed_fields=True,
                    )
                },
            )
            if repository is not None:
                try:
                    repository.record_call(
                        request=request,
                        response=response,
                        run_id=run_id,
                        status="success",
                        mode=adapter.mode,
                        external_call_id=external_call_id,
                        metadata=metadata,
                        call_id=call_id,
                    )
                except Exception as rec_exc:  # noqa: BLE001
                    logging.getLogger(__name__).warning(
                        "Failed to record successful call (call_id=%s, run_id=%s): %s",
                        call_id,
                        run_id,
                        rec_exc,
                    )

        common._emit(response.model_dump(mode="json", exclude_computed_fields=True))
    finally:
        if repository is not None:
            repository.close()
        registry.close()
