"""
Global configuration for Indic Heritage Studio v2.

MAJOR UPGRADE from v1:
  - SDXL 1.0-base + DreamShaper-XL fine-tune (vs SD 1.5 in v1)
  - Stable Video Diffusion (vs AnimateDiff in v1)
  - IP-Adapter XL for high-fidelity style transfer
  - Per-style LoRA fine-tunes (loaded on demand)
  - ControlNet (Canny / Depth / OpenPose) for composition control
  - Multi-GPU batch processing across all 8 GPUs
  - 1024×1024 native resolution (vs 512×512 in v1)

Built for 8 × NVIDIA 80GB GPUs in dev mode, with the same code running
on AMD Radeon Cloud for the final rule-compliant demo.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional

from dotenv import load_dotenv

load_dotenv()


@dataclass(frozen=True)
class Settings:
    # ------------------------------------------------------------------
    # Paths
    # ------------------------------------------------------------------
    project_root: Path = Path(__file__).resolve().parent.parent
    assets_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent / "assets")
    styles_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent / "assets" / "styles")
    outputs_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent / "outputs")
    examples_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent / "examples")
    lora_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent / "assets" / "loras")
    dataset_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent / "assets" / "datasets")

    # ------------------------------------------------------------------
    # HuggingFace model IDs — v2 upgraded stack
    # ------------------------------------------------------------------
    # Text-to-Image: SDXL 1.0 base + DreamShaper-XL fine-tune
    # DreamShaper-XL gives a richer, more artistic baseline that responds
    # well to heritage-style LoRA overlays.
    t2i_model_id: str = "lykon/dreamshaper-xl-v2-turbo"  # turbo = 25 steps, SDXL-native
    t2i_refiner_id: Optional[str] = "stabilityai/stable-diffusion-xl-refiner-1.0"

    # SDXL inpainting checkpoint (same family for consistency)
    inpaint_model_id: str = "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"

    # IP-Adapter XL — high-fidelity image-conditioned style transfer
    # IMPORTANT: use the BASE ViT-H variant (not the "plus" variant)
    # because "plus" has a 1280-dim projection that doesn't match SDXL's
    # 1664-dim image encoder. The base version is dimension-matched.
    ip_adapter_model_id: str = "h94/IP-Adapter"
    ip_adapter_subfolder: str = "sdxl_models"
    ip_adapter_weight_name: str = "ip-adapter_sdxl_vit-h.safetensors"
    ip_adapter_plus_weight_name: str = "ip-adapter-plus_sdxl_vit-h.safetensors"
    image_encoder_id: str = "h94/IP-Adapter/sdxl_models/image_encoder"

    # Stable Video Diffusion (SVD) — image-to-video, 14-25 frames, 1024×576
    svd_model_id: str = "stabilityai/stable-video-diffusion-img2vid-xt-1-1"
    svd_num_frames: int = 25
    svd_fps: int = 8
    svd_motion_bucket_id: int = 127   # 1-255, higher = more motion
    svd_noise_aug_strength: float = 0.02

    # AnimateDiff-XL fallback (used if SVD runs out of VRAM or for stylized loops)
    animatediff_xl_model_id: str = "emilianJR/animatediffXL"
    animatediff_motion_module: str = "animatediff_motion_lora_sdxl14.safetensors"
    animatediff_num_frames: int = 16

    # ControlNet — SDXL-compatible checkpoints
    controlnet_canny_id: str = "diffusers/controlnet-canny-sdxl-1.0"
    controlnet_depth_id: str = "diffusers/controlnet-depth-sdxl-1.0"
    controlnet_openpose_id: str = "thibaud/controlnet-openpose-sdxl-1.0"

    # ------------------------------------------------------------------
    # Per-style LoRA paths (each LoRA file under assets/loras/<style>.safetensors)
    # ------------------------------------------------------------------
    # LoRAs are trained in Week 1 by training/train_lora.py on heritage
    # art samples collected in assets/datasets/<style>/*.jpg.
    # If a LoRA file is missing for a given style, the pipeline falls
    # back to prompt-only conditioning (no crash).
    lora_scale_default: float = 0.85
    lora_use_fallback: bool = True

    # ------------------------------------------------------------------
    # Multi-GPU configuration
    # ------------------------------------------------------------------
    # On the dev box: 8 × 80GB GPUs = 640 GB VRAM. Strategy:
    #   - GPU 0: SDXL T2I / inpainting pipeline (resident)
    #   - GPU 1: IP-Adapter XL style transfer (resident)
    #   - GPU 2: Stable Video Diffusion (resident)
    #   - GPU 3: ControlNet preprocessing (resident)
    #   - GPU 4-7: Batch parallelism (one shard each)
    # On AMD Radeon Cloud (single GPU): all pipelines share GPU 0,
    # loaded/unloaded on demand.
    multi_gpu_enabled: bool = field(default_factory=lambda: os.getenv("MULTI_GPU", "1") == "1")
    pipeline_gpu_assignment: dict = field(default_factory=lambda: {
        "t2i": 0,
        "style": 1,
        "i2v": 2,
        "controlnet": 3,
        "inpaint": 0,         # shares with T2I (same base model)
        "batch_workers": [4, 5, 6, 7],
    })
    batch_num_workers: int = field(default_factory=lambda: int(os.getenv("BATCH_WORKERS", "4")))

    # ------------------------------------------------------------------
    # AMD Model API (agent layer — free, no GPU credits)
    # ------------------------------------------------------------------
    amd_api_key: str = field(default_factory=lambda: os.getenv("AMD_MODEL_API_KEY", ""))
    amd_base_url: str = field(default_factory=lambda: os.getenv("AMD_MODEL_BASE_URL", "https://developer.amd.com.cn/radeon/api/v1"))
    amd_agent_model: str = field(default_factory=lambda: os.getenv("AMD_AGENT_MODEL", "Qwen3.6-35B-A3B"))
    amd_agent_fallback: str = field(default_factory=lambda: os.getenv("AMD_AGENT_MODEL_FALLBACK", "DeepSeek-V4-Flash"))

    # ------------------------------------------------------------------
    # Generation defaults (v2 = SDXL-native resolution + 25 steps turbo)
    # ------------------------------------------------------------------
    torch_dtype_str: str = "float16"  # bf16 is supported on A100/H100 but fp16 for AMD parity
    default_steps: int = 25           # SDXL turbo: 25 steps is high quality
    default_steps_high: int = 50      # for showcase images
    default_guidance: float = 7.0     # SDXL prefers slightly lower CFG than SD 1.5
    default_image_size: int = 1024    # SDXL native
    default_image_size_high: int = 1280  # for showcase outputs
    default_seed: int = 42

    # SVD video defaults
    svd_frames_default: int = 25
    svd_fps_default: int = 8

    # ------------------------------------------------------------------
    # UI / Server
    # ------------------------------------------------------------------
    gradio_server_name: str = field(default_factory=lambda: os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"))
    gradio_server_port: int = field(default_factory=lambda: int(os.getenv("GRADIO_SERVER_PORT", "7860")))

    # ------------------------------------------------------------------
    # Convenience properties
    # ------------------------------------------------------------------
    @property
    def torch_dtype(self):
        import torch
        return getattr(torch, self.torch_dtype_str)

    @property
    def device(self):
        """Default device for single-GPU operations."""
        import torch
        return "cuda" if torch.cuda.is_available() else "cpu"

    @property
    def device_count(self) -> int:
        """Number of available CUDA/ROCm devices."""
        try:
            import torch
            return torch.cuda.device_count() if torch.cuda.is_available() else 0
        except Exception:
            return 0

    @property
    def is_rocm(self) -> bool:
        try:
            import torch
            if not torch.cuda.is_available():
                return False
            return hasattr(torch.version, "hip") and torch.version.hip is not None
        except Exception:
            return False

    @property
    def is_cuda(self) -> bool:
        try:
            import torch
            if not torch.cuda.is_available():
                return False
            return not self.is_rocm
        except Exception:
            return False

    @property
    def gpu_name(self) -> str:
        try:
            import torch
            if torch.cuda.is_available():
                return torch.cuda.get_device_name(0)
            return "CPU"
        except Exception:
            return "unknown"

    @property
    def vram_gb_per_gpu(self) -> float:
        """VRAM (GB) on the primary GPU."""
        try:
            import torch
            if torch.cuda.is_available():
                _, total = torch.cuda.mem_get_info()
                return round(total / 1e9, 1)
        except Exception:
            pass
        return 0.0

    @property
    def total_vram_gb(self) -> float:
        """Total VRAM across all GPUs (GB)."""
        try:
            import torch
            if not torch.cuda.is_available():
                return 0.0
            total = 0.0
            for i in range(torch.cuda.device_count()):
                _, mem = torch.cuda.mem_get_info(i)
                total += mem
            return round(total / 1e9, 1)
        except Exception:
            return 0.0

    @property
    def agents_enabled(self) -> bool:
        return bool(self.amd_api_key)

    def get_pipeline_device(self, pipeline_name: str) -> str:
        """Return the cuda device for the given pipeline (multi-GPU aware)."""
        if not self.multi_gpu_enabled or self.device_count <= 1:
            return self.device
        idx = self.pipeline_gpu_assignment.get(pipeline_name, 0)
        if isinstance(idx, list):
            idx = idx[0]
        return f"cuda:{min(idx, self.device_count - 1)}"

    def get_batch_worker_devices(self) -> List[str]:
        """Return list of cuda devices for batch parallelism."""
        if not self.multi_gpu_enabled or self.device_count <= 1:
            return [self.device]
        workers = self.pipeline_gpu_assignment.get("batch_workers", [])
        if not workers:
            workers = list(range(min(self.batch_num_workers, self.device_count)))
        return [f"cuda:{min(i, self.device_count - 1)}" for i in workers]

    def lora_path_for(self, style_id: str) -> Optional[Path]:
        """Return the LoRA safetensors path for a style, if it exists."""
        p = self.lora_dir / f"{style_id}.safetensors"
        return p if p.exists() else None

    def ensure_dirs(self) -> None:
        for d in (self.outputs_dir, self.examples_dir, self.styles_dir,
                  self.lora_dir, self.dataset_dir):
            d.mkdir(parents=True, exist_ok=True)


# Singleton
settings = Settings()
settings.ensure_dirs()