import os
import uuid
from datetime import datetime, timezone
from sqlalchemy import select
from app.tasks.celery_app import celery_app
from app.core.database import async_session_factory
from app.core.config import get_settings
from app.models.generation_task import GenerationTask
from app.models.generation_point import GenerationPoint
from app.models.template import Template
from app.services.document_processor import html_to_docx_bytes
from app.services.ai_adapter import call_ai_model
from app.services.ref_parser import parse_reference_file
from app.services.file_storage import get_storage_dir, RESULTS_DIR


async def _generate_document(task_id: str) -> None:
    settings = get_settings()

    async with async_session_factory() as db:
        task_result = await db.execute(select(GenerationTask).where(GenerationTask.id == task_id))
        task = task_result.scalar_one_or_none()
        if not task:
            return

        task.status = "processing"
        await db.commit()

        template_result = await db.execute(select(Template).where(Template.id == task.template_id))
        template = template_result.scalar_one_or_none()
        if not template:
            task.status = "failed"
            task.error_msg = "模板不存在"
            task.finished_at = datetime.now(timezone.utc)
            await db.commit()
            return

        html_content = template.html_content or ""

        points_result = await db.execute(
            select(GenerationPoint)
            .where(GenerationPoint.template_id == task.template_id)
            .order_by(GenerationPoint.order.asc())
        )
        points = points_result.scalars().all()

        try:
            total = len(points)
            for idx, point in enumerate(points):
                ref_content = None
                if point.ref_file_path:
                    ref_content = await parse_reference_file(point.ref_file_path)

                model_config = {"provider": "custom", "endpoint": "", "api_key": "", "extra_params": {}}
                if point.model_id:
                    from app.models.ai_model import AIModel
                    model_result = await db.execute(select(AIModel).where(AIModel.id == point.model_id))
                    model = model_result.scalar_one_or_none()
                    if model:
                        model_config = {
                            "provider": model.provider,
                            "endpoint": model.endpoint,
                            "api_key": model.api_key,
                            "extra_params": model.extra_params,
                        }

                ai_result = await call_ai_model(model_config, point.prompt, ref_content)

                position = point.position
                start = position.get("start", 0)
                end = position.get("end", 0)
                if 0 <= start < end <= len(html_content):
                    html_content = html_content[:start] + ai_result + html_content[end:]

                task.status = "processing"
                await db.commit()

            result_dir = get_storage_dir(RESULTS_DIR)
            result_path = os.path.join(result_dir, f"{task_id}.docx")
            docx_bytes = html_to_docx_bytes(html_content)
            with open(result_path, "wb") as f:
                f.write(docx_bytes)

            task.status = "done"
            task.result_file_path = result_path
            task.finished_at = datetime.now(timezone.utc)
            await db.commit()

        except Exception as e:
            task.status = "failed"
            task.error_msg = str(e)
            task.finished_at = datetime.now(timezone.utc)
            await db.commit()


@celery_app.task(bind=True, name="generate_document")
def generate_document(self, task_id: str) -> dict:
    import asyncio
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(_generate_document(task_id))
        return {"status": "done", "task_id": task_id}
    except Exception as e:
        return {"status": "failed", "task_id": task_id, "error": str(e)}
    finally:
        loop.close()