from __future__ import annotations

import socket

from app.configs.settings import get_settings
from app.models.common import utc_now
from app.models.domain import JobProgress
from app.services.comfyui_service import ComfyUIService
from app.services.container import AppContainer
from app.services.huggingface_service import HuggingFaceService
from app.services.workflow_service import WorkflowService
from app.utils.enums import CreditReason, JobStatus


class WorkerRuntime:
    def __init__(self) -> None:
        self._settings = get_settings()
        self._container = AppContainer(self._settings)
        self._comfyui = ComfyUIService(self._settings)
        self._workflow = WorkflowService()
        self._huggingface = (
            HuggingFaceService(self._settings)
            if self._settings.features.enable_model_auto_sync
            else None
        )

    async def startup(self) -> None:
        await self._container.startup()

    async def shutdown(self) -> None:
        await self._container.shutdown()

    async def process_job(self, job_id: str, celery_task_id: str) -> None:
        if self._container.repositories is None or self._container.services is None:
            raise RuntimeError("Worker container did not initialize.")
        repositories = self._container.repositories
        services = self._container.services

        job = await repositories.jobs.get_by_id(job_id)
        if job is None:
            raise ValueError(f"Job {job_id} not found.")
        user = await repositories.users.get_by_id(job.user_id)
        if user is None:
            raise ValueError(f"User {job.user_id} not found.")

        await repositories.jobs.update_status(
            job_id,
            status=JobStatus.RUNNING,
            progress=JobProgress(percent=2, stage="initializing", message="Preparing worker"),
            fields={
                "celery_task_id": celery_task_id,
                "assigned_worker": socket.gethostname(),
                "started_at": utc_now(),
            },
        )

        model = services.models.get_model(job.options.model_key)
        if self._huggingface is not None:
            await self._huggingface.ensure_model_available(model)

        source_filename: str | None = None
        if job.options.source_image_id:
            source_image = await repositories.images.get_by_id(job.options.source_image_id)
            if source_image is None:
                raise ValueError("Source image not found.")
            source_bytes = await services.storage.read(source_image.storage_key)
            source_filename = f"{source_image.id}.png"
            await self._comfyui.upload_image(source_filename, source_bytes, source_image.mime_type)

        workflow = self._workflow.build(job, model, source_filename=source_filename)

        async def on_progress(percent: int, stage: str) -> None:
            await repositories.jobs.update_status(
                job_id,
                status=JobStatus.RUNNING,
                progress=JobProgress(percent=percent, stage=stage, message=stage.title()),
            )

        execution = await self._comfyui.execute_workflow(workflow, progress_cb=on_progress)
        image_ids: list[str] = []
        for image_bytes in execution.images:
            created_image = await services.images.store_generated_image(user.id or "", job_id, image_bytes)
            image_ids.append(created_image.id or "")

        await repositories.jobs.attach_images(job_id, image_ids)
        await repositories.jobs.update_status(
            job_id,
            status=JobStatus.SUCCEEDED,
            progress=JobProgress(percent=100, stage="completed", message="Completed"),
            fields={
                "finished_at": utc_now(),
                "comfy_prompt_id": execution.prompt_id,
                "image_ids": image_ids,
            },
        )

    async def mark_retrying(self, job_id: str, error_message: str) -> None:
        if self._container.repositories is None:
            return
        await self._container.repositories.jobs.update_status(
            job_id,
            status=JobStatus.RETRYING,
            progress=JobProgress(percent=0, stage="retrying", message="Retrying"),
            error_message=error_message,
        )

    async def mark_failed(self, job_id: str, error_message: str) -> None:
        if self._container.repositories is None or self._container.services is None:
            return
        repositories = self._container.repositories
        services = self._container.services
        job = await repositories.jobs.get_by_id(job_id)
        if job is None:
            return
        user = await repositories.users.get_by_id(job.user_id)
        if user is not None and job.charged_credits > 0:
            await services.credits.grant(
                user=user,
                amount=job.charged_credits,
                reason=CreditReason.REFUND,
                note=f"Refund for failed job {job_id}",
            )
        await repositories.jobs.update_status(
            job_id,
            status=JobStatus.FAILED,
            progress=JobProgress(percent=0, stage="failed", message="Failed"),
            error_message=error_message,
            fields={"finished_at": utc_now()},
        )

    async def cleanup_expired_images(self) -> int:
        if self._container.repositories is None or self._container.services is None:
            return 0
        repositories = self._container.repositories
        services = self._container.services
        expired_images = await repositories.images.list_expired(limit=500)
        deleted = 0
        for image in expired_images:
            await services.storage.delete(image.storage_key)
            if image.id:
                await repositories.images.delete(image.id)
            deleted += 1
        return deleted