"""Thumbnail rendering tasks.

Covers:
- render_step_thumbnail  — render thumbnail for a freshly-processed STEP file
- regenerate_thumbnail   — re-render thumbnail with updated per-part colours
"""
import logging
from contextlib import contextmanager
from pathlib import Path

from app.tasks.celery_app import celery_app
from app.core.task_logs import log_task_event
from app.core.pipeline_logger import PipelineLogger

logger = logging.getLogger(__name__)

# Maximum samples for thumbnail renders (512x512).
# Full-resolution renders use 256+ samples; thumbnails don't need more than 64.
_THUMBNAIL_SAMPLE_CAP = 64


def _resolve_thumbnail_render_context(session, cad) -> dict[str, object]:
    """Reuse workflow material/USD resolution for CAD thumbnails when possible."""
    context: dict[str, object] = {}
    if not cad:
        return context

    parsed_objects = cad.parsed_objects if isinstance(cad.parsed_objects, dict) else {}
    raw_part_names = parsed_objects.get("objects") if isinstance(parsed_objects, dict) else None
    if isinstance(raw_part_names, list):
        part_names_ordered = [
            str(part_name).strip()
            for part_name in raw_part_names
            if isinstance(part_name, str) and part_name.strip()
        ]
        if part_names_ordered:
            context["part_names_ordered"] = part_names_ordered

    try:
        from sqlalchemy import select

        from app.core.render_paths import resolve_result_path
        from app.domains.media.models import MediaAsset, MediaAssetType
        from app.domains.products.models import Product
        from app.domains.rendering.workflow_runtime_services import (
            _build_effective_material_lookup,
            _usd_master_refresh_reason,
        )
        from app.services.material_service import resolve_material_map
        from app.services.template_service import get_material_library_path_for_session

        product = session.execute(
            select(Product)
            .where(Product.cad_file_id == cad.id)
            .order_by(Product.is_active.desc(), Product.updated_at.desc(), Product.created_at.desc())
            .limit(1)
        ).scalar_one_or_none()

        material_library_path = get_material_library_path_for_session(session)
        materials_source = product.cad_part_materials or [] if product else []
        raw_material_map = _build_effective_material_lookup(cad, materials_source)
        if material_library_path and raw_material_map:
            material_map = resolve_material_map(raw_material_map)
            if material_map:
                context["material_library_path"] = material_library_path
                context["material_map"] = material_map

        usd_asset = session.execute(
            select(MediaAsset)
            .where(
                MediaAsset.cad_file_id == cad.id,
                MediaAsset.asset_type == MediaAssetType.usd_master,
            )
            .order_by(MediaAsset.created_at.desc())
            .limit(1)
        ).scalar_one_or_none()
        if usd_asset:
            usd_path = resolve_result_path(usd_asset.storage_key)
            refresh_reason = _usd_master_refresh_reason(
                cad,
                usd_asset=usd_asset,
                usd_render_path=usd_path,
            )
            if refresh_reason is None and usd_path and usd_path.exists():
                context["usd_path"] = usd_path
    except Exception:
        logger.exception("Failed to resolve thumbnail render context for cad %s", getattr(cad, "id", None))

    return context


def _render_thumbnail_core(
    *,
    cad_file_id: str,
    workflow_run_id: str | None = None,
    workflow_node_id: str | None = None,
    renderer: str | None = None,
    render_engine: str | None = None,
    samples: int | None = None,
    width: int | None = None,
    height: int | None = None,
    transparent_bg: bool | None = None,
    include_postprocess: bool,
    queue_legacy_glb_follow_up: bool,
) -> None:
    """Render a CAD thumbnail with optional legacy post-processing."""
    pl = PipelineLogger(task_id=None)
    pl.step_start("render_step_thumbnail", {"cad_file_id": cad_file_id})
    logger.info("Rendering thumbnail for CAD file: %s", cad_file_id)

    from app.core.tenant_context import resolve_tenant_id_for_cad

    tenant_id = resolve_tenant_id_for_cad(cad_file_id)

    try:
        from app.models.cad_file import CadFile
        from app.domains.products.cache_service import compute_step_hash

        with _pipeline_session(tenant_id) as session:
            cad = session.get(CadFile, cad_file_id)
            if cad and cad.stored_path and not cad.step_file_hash:
                cad.step_file_hash = compute_step_hash(cad.stored_path)
                session.commit()
                logger.info("Saved step_file_hash for %s: %s…", cad_file_id, cad.step_file_hash[:12])
    except Exception:
        logger.warning("step_file_hash computation failed for %s (non-fatal)", cad_file_id)

    render_context: dict[str, object] = {}
    try:
        from app.models.cad_file import CadFile

        with _pipeline_session(tenant_id) as session:
            cad = session.get(CadFile, cad_file_id)
            render_context = _resolve_thumbnail_render_context(session, cad)
    except Exception:
        logger.warning("thumbnail render context resolution failed for %s; using fallback render path", cad_file_id)

    try:
        from app.services.step_processor import regenerate_cad_thumbnail

        pl.info("render_step_thumbnail", "Calling regenerate_cad_thumbnail")
        with _capped_thumbnail_samples():
            success = regenerate_cad_thumbnail(
                cad_file_id,
                part_colors={},
                renderer=renderer,
                render_engine=render_engine,
                samples=samples,
                width=width,
                height=height,
                transparent_bg=transparent_bg,
                **render_context,
            )
        if not success:
            raise RuntimeError("regenerate_cad_thumbnail returned False")
    except Exception as exc:
        pl.step_error("render_step_thumbnail", f"Thumbnail render failed: {exc}", exc)
        logger.error("Thumbnail render failed for %s: %s", cad_file_id, exc)
        raise

    resolved_tenant_id: str | None = None
    if include_postprocess:
        try:
            from app.models.cad_file import CadFile
            from app.domains.rendering.workflow_runtime_services import resolve_cad_bbox

            with _pipeline_session(tenant_id) as session:
                cad = session.get(CadFile, cad_file_id)
                if not cad:
                    logger.warning("CadFile %s not found in post-render phase", cad_file_id)
                else:
                    step_path = cad.stored_path
                    attrs = cad.mesh_attributes or {}

                    if step_path and not attrs.get("dimensions_mm"):
                        step_file = Path(step_path)
                        glb_path = step_file.parent / f"{step_file.stem}_thumbnail.glb"
                        bbox_data = resolve_cad_bbox(step_path, glb_path=str(glb_path)).bbox_data
                        if bbox_data:
                            cad.mesh_attributes = {**attrs, **bbox_data}
                            attrs = cad.mesh_attributes
                            dims = bbox_data["dimensions_mm"]
                            logger.info(
                                "bbox for %s: %s×%s×%s mm",
                                cad_file_id,
                                dims["x"],
                                dims["y"],
                                dims["z"],
                            )

                    if step_path and "sharp_edge_pairs" not in attrs:
                        try:
                            from app.services.step_processor import extract_mesh_edge_data

                            edge_data = extract_mesh_edge_data(step_path)
                            if edge_data:
                                cad.mesh_attributes = {**attrs, **edge_data}
                                n_pairs = len(edge_data.get("sharp_edge_pairs", []))
                                logger.info(
                                    "Sharp edge data extracted for %s: %s sharp edges",
                                    cad_file_id,
                                    n_pairs,
                                )
                        except Exception:
                            logger.exception(
                                "Sharp edge extraction failed for %s (non-fatal)",
                                cad_file_id,
                            )

                    session.commit()
                    resolved_tenant_id = str(cad.tenant_id) if cad.tenant_id else None
        except Exception:
            logger.exception("Post-render processing failed for %s (non-fatal)", cad_file_id)

        try:
            from app.domains.pipeline.tasks.extract_metadata import _auto_populate_materials_for_cad

            _auto_populate_materials_for_cad(cad_file_id, tenant_id=tenant_id)
        except Exception:
            logger.exception(
                "Auto material population failed for cad_file %s (non-fatal)",
                cad_file_id,
            )

        try:
            if resolved_tenant_id:
                from app.core.websocket import publish_event_sync

                publish_event_sync(
                    resolved_tenant_id,
                    {
                        "type": "cad_processing_complete",
                        "cad_file_id": cad_file_id,
                        "status": "completed",
                    },
                )
        except Exception:
            logger.debug("WebSocket publish for CAD complete skipped (non-fatal)")

    if queue_legacy_glb_follow_up:
        try:
            from app.domains.pipeline.tasks.export_glb import generate_gltf_geometry_task

            generate_gltf_geometry_task.delay(cad_file_id)
            pl.info("render_step_thumbnail", f"Queued generate_gltf_geometry_task for {cad_file_id}")
        except Exception:
            logger.debug("Could not queue generate_gltf_geometry_task (non-fatal)")

    pl.step_done("render_step_thumbnail")
    try:
        from app.domains.rendering.tasks import _update_workflow_run_status

        _update_workflow_run_status(
            cad_file_id,
            "completed",
            workflow_run_id=workflow_run_id,
            workflow_node_id=workflow_node_id,
        )
    except Exception:
        logger.exception("Failed to update workflow state for thumbnail render %s", cad_file_id)


@contextmanager
def _capped_thumbnail_samples():
    """Temporarily cap render samples for thumbnail renders.

    Thumbnails are 512x512 — using 256 Cycles samples is wasteful.
    This patches _get_all_settings in step_processor to cap samples
    at _THUMBNAIL_SAMPLE_CAP for the duration of the thumbnail render.
    """
    import app.services.step_processor as _sp
    _original = _sp._get_all_settings

    def _patched() -> dict[str, str]:
        settings = _original()
        for key in ("blender_cycles_samples", "blender_eevee_samples"):
            try:
                val = int(settings.get(key, "256"))
                if val > _THUMBNAIL_SAMPLE_CAP:
                    logger.info("Capping thumbnail %s: %d -> %d", key, val, _THUMBNAIL_SAMPLE_CAP)
                    settings[key] = str(_THUMBNAIL_SAMPLE_CAP)
            except (ValueError, TypeError):
                pass
        return settings

    _sp._get_all_settings = _patched
    try:
        yield
    finally:
        _sp._get_all_settings = _original


@contextmanager
def _pipeline_session(tenant_id: str | None = None):
    """Single DB engine + session for the entire task lifetime.

    Replaces the previous pattern of creating 3-7 separate create_engine()
    + Session() pairs per task invocation. Each create_engine() spins up a
    new connection pool, wasting ~50ms + one PG connection per call.
    """
    from sqlalchemy import create_engine
    from sqlalchemy.orm import Session
    from app.config import settings as app_settings
    from app.core.tenant_context import set_tenant_context_sync

    sync_url = app_settings.database_url.replace("+asyncpg", "")
    engine = create_engine(sync_url)
    try:
        with Session(engine) as session:
            set_tenant_context_sync(session, tenant_id)
            yield session
    finally:
        engine.dispose()


@celery_app.task(bind=True, name="app.tasks.step_tasks.render_step_thumbnail", queue="asset_pipeline")
def render_step_thumbnail(
    self,
    cad_file_id: str,
    workflow_run_id: str | None = None,
    workflow_node_id: str | None = None,
    renderer: str | None = None,
    render_engine: str | None = None,
    samples: int | None = None,
    width: int | None = None,
    height: int | None = None,
    transparent_bg: bool | None = None,
    **_: object,
):
    """Render the thumbnail for a freshly-processed STEP file.

    Runs on the dedicated asset_pipeline queue (concurrency=1) so the
    blender-renderer service is never overwhelmed by concurrent requests.
    On success, also auto-populates materials and marks the CadFile as completed.
    """
    try:
        _render_thumbnail_core(
            cad_file_id=cad_file_id,
            workflow_run_id=workflow_run_id,
            workflow_node_id=workflow_node_id,
            renderer=renderer,
            render_engine=render_engine,
            samples=samples,
            width=width,
            height=height,
            transparent_bg=transparent_bg,
            include_postprocess=True,
            queue_legacy_glb_follow_up=workflow_run_id is None,
        )
    except Exception as exc:
        raise self.retry(exc=exc, countdown=30, max_retries=2)


@celery_app.task(bind=True, name="app.tasks.step_tasks.render_graph_thumbnail", queue="asset_pipeline")
def render_graph_thumbnail(
    self,
    cad_file_id: str,
    workflow_run_id: str | None = None,
    workflow_node_id: str | None = None,
    renderer: str | None = None,
    render_engine: str | None = None,
    samples: int | None = None,
    width: int | None = None,
    height: int | None = None,
    transparent_bg: bool | None = None,
    **_: object,
):
    """Render a CAD thumbnail for graph workflows without legacy follow-up side effects."""
    try:
        _render_thumbnail_core(
            cad_file_id=cad_file_id,
            workflow_run_id=workflow_run_id,
            workflow_node_id=workflow_node_id,
            renderer=renderer,
            render_engine=render_engine,
            samples=samples,
            width=width,
            height=height,
            transparent_bg=transparent_bg,
            include_postprocess=False,
            queue_legacy_glb_follow_up=False,
        )
    except Exception as exc:
        raise self.retry(exc=exc, countdown=30, max_retries=2)


@celery_app.task(bind=True, name="app.tasks.step_tasks.regenerate_thumbnail", queue="asset_pipeline")
def regenerate_thumbnail(
    self,
    cad_file_id: str,
    part_colors: dict,
    renderer: str | None = None,
    render_engine: str | None = None,
    samples: int | None = None,
    width: int | None = None,
    height: int | None = None,
    transparent_bg: bool | None = None,
):
    """Regenerate thumbnail with per-part colours."""
    pl = PipelineLogger(task_id=self.request.id)
    pl.step_start("regenerate_thumbnail", {"cad_file_id": cad_file_id})
    logger.info(f"Regenerating thumbnail for CAD file: {cad_file_id}")

    # Resolve and log tenant context at task start (required for RLS)
    from app.core.tenant_context import resolve_tenant_id_for_cad
    _tenant_id = resolve_tenant_id_for_cad(cad_file_id)

    try:
        from app.services.step_processor import MissingCadResourceError, regenerate_cad_thumbnail

        render_context: dict[str, object] = {}
        try:
            from app.models.cad_file import CadFile

            with _pipeline_session(_tenant_id) as session:
                cad = session.get(CadFile, cad_file_id)
                render_context = _resolve_thumbnail_render_context(session, cad)
        except Exception:
            logger.warning(
                "thumbnail render context resolution failed for %s during regeneration; using fallback render path",
                cad_file_id,
            )

        with _capped_thumbnail_samples():
            success = regenerate_cad_thumbnail(
                cad_file_id,
                part_colors,
                renderer=renderer,
                render_engine=render_engine,
                samples=samples,
                width=width,
                height=height,
                transparent_bg=transparent_bg,
                **render_context,
            )
        if not success:
            raise RuntimeError("regenerate_cad_thumbnail returned False")
    except MissingCadResourceError as exc:
        pl.warning("regenerate_thumbnail", f"Skipping stale thumbnail regeneration: {exc}")
        logger.warning("Skipping thumbnail regeneration for %s: %s", cad_file_id, exc)
        pl.step_done("regenerate_thumbnail")
        return
    except Exception as exc:
        pl.step_error("regenerate_thumbnail", f"Thumbnail regeneration failed: {exc}", exc)
        logger.error(f"Thumbnail regeneration failed for {cad_file_id}: {exc}")
        raise self.retry(exc=exc, countdown=30, max_retries=2)
    pl.step_done("regenerate_thumbnail")