from datetime import datetime
from fastapi import HTTPException
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from models.ollama_model_status import OllamaModelStatus
from models.sql.ollama_pull_status import OllamaPullStatus
from services.database import async_session_maker
from utils.ollama import pull_ollama_model


async def pull_ollama_model_background_task(model: str):
    saved_model_status = OllamaModelStatus(
        name=model,
        status="pulling",
        done=False,
    )
    log_event_count = 0

    async with async_session_maker() as session:
        try:
            async for event in pull_ollama_model(model):
                log_event_count += 1
                if log_event_count != 1 and log_event_count % 20 != 0:
                    continue

                if "completed" in event:
                    saved_model_status.downloaded = event["completed"]

                if not saved_model_status.size and "total" in event:
                    saved_model_status.size = event["total"]

                if "status" in event:
                    saved_model_status.status = event["status"]

                    await upsert_ollama_pull_status(session, model, saved_model_status)

        except Exception as e:
            saved_model_status.status = "error"
            saved_model_status.done = True
            await upsert_ollama_pull_status(session, model, saved_model_status)
            raise HTTPException(
                status_code=500,
                detail=f"Failed to pull model: {e}",
            )

        saved_model_status.done = True
        saved_model_status.status = "pulled"
        saved_model_status.downloaded = saved_model_status.size

        await upsert_ollama_pull_status(session, model, saved_model_status)


async def upsert_ollama_pull_status(
    session: AsyncSession, model: str, model_status: OllamaModelStatus
):
    stmt = select(OllamaPullStatus).where(OllamaPullStatus.id == model)
    result = await session.execute(stmt)
    existing_record = result.scalar_one_or_none()

    if existing_record:
        existing_record.status = model_status.model_dump(mode="json")
        existing_record.last_updated = datetime.now()
    else:
        new_record = OllamaPullStatus(
            id=model,
            status=model_status.model_dump(mode="json"),
            last_updated=datetime.now(),
        )
        session.add(new_record)

    await session.commit()
    await session.flush()
