import dataclasses
import json
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import override

from pydantic_ai import RunContext
from pydantic_ai.messages import ModelMessage, ModelResponse
from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
from pydantic_ai.settings import ModelSettings


def _default(obj: object) -> object:
    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
        return dataclasses.asdict(obj)
    return str(obj)


class LoggingModel(Model):
    _inner: Model
    _log_dir: Path

    def __init__(self, inner: Model, log_dir: Path) -> None:
        super().__init__()
        self._inner = inner
        self._log_dir = log_dir
        log_dir.mkdir(parents=True, exist_ok=True)

    @property
    @override
    def model_name(self) -> str:
        return self._inner.model_name

    @property
    @override
    def system(self) -> str:
        return self._inner.system

    def _save(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
    ) -> None:
        timestamp = datetime.now().strftime("%Y%m%dT%H%M%S%f")
        payload: dict[str, object] = {
            "timestamp": datetime.now().isoformat(),
            "messages": [dataclasses.asdict(m) for m in messages],
            "model_settings": dict(model_settings) if model_settings else None,
            "model_request_parameters": dataclasses.asdict(model_request_parameters),
        }
        log_path = self._log_dir / f"request_{timestamp}.json"
        log_path.write_text(
            json.dumps(payload, indent=2, ensure_ascii=False, default=_default),
            encoding="utf-8",
        )

    @override
    async def request(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
    ) -> ModelResponse:
        self._save(messages, model_settings, model_request_parameters)
        return await self._inner.request(messages, model_settings, model_request_parameters)

    @asynccontextmanager
    @override
    async def request_stream(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
        run_context: RunContext[object] | None = None,
    ) -> AsyncIterator[StreamedResponse]:
        self._save(messages, model_settings, model_request_parameters)
        async with self._inner.request_stream(
            messages, model_settings, model_request_parameters, run_context
        ) as stream:
            yield stream
