import json
import logging
import time
import traceback
from typing import AsyncIterator, AsyncGenerator
from contextlib import asynccontextmanager

import litellm
from app.core.config import DEFAULT_MODEL, DEBUG

litellm.drop_params = True
if DEBUG:
    litellm._turn_on_debug()

logger = logging.getLogger(__name__)


@asynccontextmanager
async def _log_context(
    model: str,
    messages: list[dict],
    session_id: str | None = None,
) -> AsyncGenerator[dict, None]:
    """LLM呼び出しの前後でログデータを収集するコンテキストマネージャー。"""
    log: dict = {
        "session_id": session_id,
        "model": model,
        "request_json": json.dumps(messages, ensure_ascii=False),
        "response_text": None,
        "error": None,
        "latency_ms": None,
    }
    start = time.perf_counter()
    try:
        yield log
    except Exception as e:
        log["error"] = str(e)
        log["latency_ms"] = round((time.perf_counter() - start) * 1000, 1)
        logger.error("LLM call failed [model=%s]: %s\n%s", log["model"], e, traceback.format_exc())
        raise
    finally:
        if log["latency_ms"] is None:
            log["latency_ms"] = round((time.perf_counter() - start) * 1000, 1)


def _save_log(log: dict) -> None:
    from app.db.database import SessionLocal
    from app.db.models import ApiLog
    db = SessionLocal()
    try:
        db.add(ApiLog(**log))
        db.commit()
    finally:
        db.close()


async def stream_chat(
    messages: list[dict],
    model: str | None = None,
    session_id: str | None = None,
) -> AsyncIterator[str]:
    used_model = model or DEFAULT_MODEL
    async with _log_context(used_model, messages, session_id) as log:
        response = await litellm.acompletion(
            model=used_model,
            messages=messages,
            stream=True,
        )
        chunks: list[str] = []
        async for chunk in response:
            delta = chunk.choices[0].delta.content
            if delta:
                chunks.append(delta)
                yield delta
        log["response_text"] = "".join(chunks)
    _save_log(log)


async def complete(
    messages: list[dict],
    model: str | None = None,
    session_id: str | None = None,
) -> str:
    used_model = model or DEFAULT_MODEL
    async with _log_context(used_model, messages, session_id) as log:
        response = await litellm.acompletion(
            model=used_model,
            messages=messages,
            stream=False,
        )
        result = response.choices[0].message.content or ""
        log["response_text"] = result
    _save_log(log)
    return result