from datetime import datetime
from typing import Optional

from app.core.database import get_db
from app.models.parse_run import DocParseRun, ParseRunStatus


def _serialize(doc: dict) -> dict:
    if doc and "_id" in doc:
        doc["id"] = str(doc.pop("_id"))
    return doc


async def create_run(run: DocParseRun) -> str:
    db = get_db()
    result = await db.doc_parse_runs.insert_one(run.model_dump(exclude={"id"}))
    return str(result.inserted_id)


async def get_run_by_id(run_id: str) -> Optional[DocParseRun]:
    db = get_db()
    doc = await db.doc_parse_runs.find_one({"run_id": run_id})
    return DocParseRun(**_serialize(doc)) if doc else None


async def update_run_status(
    run_id: str,
    status: str,
    chunks_created: int = 0,
    error: Optional[str] = None,
) -> None:
    db = get_db()
    update: dict = {"status": status, "chunks_created": chunks_created}
    if status == ParseRunStatus.PROCESSING:
        update["started_at"] = datetime.utcnow()
    if status in (ParseRunStatus.COMPLETED, ParseRunStatus.FAILED):
        update["completed_at"] = datetime.utcnow()
    if error:
        update["error"] = error
    await db.doc_parse_runs.update_one({"run_id": run_id}, {"$set": update})