from __future__ import annotations

import json
import traceback
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

from .config import WorkflowConfig, ensure_job_dirs
from .stages import (
    materialize_artifacts,
    run_rig_facefusion_stage,
    run_shape_stage,
    run_texture_stage,
)


class UnifiedWorkflow:
    def __init__(
        self, config: WorkflowConfig, output_dir: Path, seed: int, tex_seed: int
    ):
        self.config = config
        self.output_dir = output_dir
        self.seed = int(seed)
        self.tex_seed = int(tex_seed)
        self.artifacts_dir, self.logs_dir = ensure_job_dirs(output_dir)

    def run(self, image_path: Path) -> dict[str, Any]:
        started = datetime.now(timezone.utc)
        manifest: dict[str, Any] = {
            "workflow": "meshforge-v2-unified",
            "started_at": started.isoformat(),
            "input": {"image": str(image_path)},
            "config": asdict(self.config),
            "seed": self.seed,
            "tex_seed": self.tex_seed,
            "stages": [],
            "artifacts": {},
            "status": "running",
        }

        all_outputs: dict[str, Any] = {}

        try:
            image_array = np.array(Image.open(image_path).convert("RGB"))

            shape = run_shape_stage(image_array, self.config, self.seed)
            all_outputs.update(shape)
            manifest["stages"].append(
                {"name": "shape", "status": "ok", "detail": shape["status"]}
            )

            texture = run_texture_stage(
                shape["shape_glb"], image_array, self.config, self.tex_seed
            )
            all_outputs.update(texture)
            manifest["stages"].append(
                {"name": "texture", "status": "ok", "detail": texture["status"]}
            )

            fused = run_rig_facefusion_stage(
                texture["textured_glb"], image_array, self.config
            )
            all_outputs.update(fused)
            manifest["stages"].append(
                {
                    "name": "rig_pshuman_fusion",
                    "status": "ok",
                    "detail": fused["status"],
                }
            )

            manifest["artifacts"] = materialize_artifacts(
                all_outputs, self.artifacts_dir
            )
            manifest["status"] = "completed"
        except Exception as exc:
            manifest["status"] = "failed"
            manifest["error"] = {
                "message": str(exc),
                "traceback": traceback.format_exc(),
            }

        manifest["finished_at"] = datetime.now(timezone.utc).isoformat()
        self._write_manifest(manifest)
        return manifest

    def _write_manifest(self, manifest: dict[str, Any]) -> None:
        out = self.output_dir / "manifest.json"
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(manifest, indent=2), encoding="utf-8")