import uuid

from sqlalchemy import select

from app.tasks.celery_app import celery_app
from app.tasks.transcription import SyncSession
from app.models.protocol import Protocol, ProtocolStatus
from app.models.transcription import Transcription
from app.services.protocol_generator import (
    generate_protocol_text,
    get_default_template,
    get_template_by_id,
)


@celery_app.task(bind=True, name="app.tasks.protocol.generate_protocol")
def generate_protocol(self, transcription_id: str, template_id: str | None = None):
    trans_uuid = uuid.UUID(transcription_id)

    with SyncSession() as db:
        transcription = db.execute(
            select(Transcription).where(Transcription.id == trans_uuid)
        ).scalar_one_or_none()

        if not transcription:
            return {"error": "Transcription not found"}

        if template_id:
            template = get_template_by_id(db, uuid.UUID(template_id))
        else:
            template = get_default_template(db)

        if not template:
            return {"error": "No template found"}

        protocol = Protocol(
            transcription_id=trans_uuid,
            template_id=template.id,
            llm_model="",
            status=ProtocolStatus.generating,
        )
        db.add(protocol)
        db.commit()
        db.refresh(protocol)

        try:
            result = generate_protocol_text(transcription.text, template)

            protocol.content = result["content"]
            protocol.raw_text = result["raw_text"]
            protocol.llm_model = result["model"]
            protocol.status = ProtocolStatus.done
            db.commit()

            return {"protocol_id": str(protocol.id), "status": "done"}

        except Exception as e:
            protocol.status = ProtocolStatus.error
            db.commit()
            raise self.retry(exc=e, max_retries=2, countdown=60)