"""
Master pipeline combining all models.
"""

import cv2
import numpy as np
from pathlib import Path

from app.config import PipelineConfig, ModelPaths, ContentType, LineArtStyle, Quality
from app.models.registry import ModelRegistry
from app.models.line_art import MultiLineArtExtractor
from app.models.sam_segmenter import SAMSegmenter, HybridSegmenter
from app.models.depth_estimator import DepthEstimator
from app.models.super_resolution import SuperResolution
from app.models.background import BackgroundRemover
from app.models.face_parser import FaceParser
from app.models.generator import TemplateGenerator
from app.processing.quantizer import (
    quantize_colors,
    merge_line_art_with_regions,
    merge_small_regions_into_neighbors,
    smooth_label_map,
    estimate_k,
)
from app.processing.utils import get_adjacency_matrix, get_pole_of_inaccessibility
from app.processing.vectorizer import generate_svg
from app.processing.difficulty import compute_region_difficulty


class ColorByNumberPipeline:
    """Full multi-model pipeline."""

    def __init__(
        self,
        model_paths: ModelPaths | None = None,
        generation_api_key: str = "",
        generation_provider: str = "replicate",
    ):
        self.paths = model_paths or ModelPaths()
        self.registry = ModelRegistry(self.paths)

        self.line_art = MultiLineArtExtractor(self.registry)
        self.sam = SAMSegmenter(self.registry)
        self.hybrid_seg = HybridSegmenter(self.sam)
        self.depth = DepthEstimator(self.registry)
        self.super_res = SuperResolution(self.registry)
        self.bg_remover = BackgroundRemover(self.registry)
        self.face_parser = FaceParser(self.registry)
        self.generator = TemplateGenerator(
            api_key=generation_api_key,
            provider=generation_provider,
        )

    def _detect_content_type(self, image: np.ndarray) -> ContentType:
        if self.face_parser.available and self.face_parser.is_portrait(image):
            return ContentType.PORTRAIT

        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        avg_sat = hsv[:, :, 1].mean()
        sat_std = hsv[:, :, 1].std()

        if avg_sat > 100 and sat_std < 50:
            return ContentType.ANIME

        h, w = image.shape[:2]
        if w / h > 1.3:
            blue_mask = (hsv[:, :, 0] > 90) & (hsv[:, :, 0] < 130)
            green_mask = (hsv[:, :, 0] > 35) & (hsv[:, :, 0] < 85)
            nature_ratio = (blue_mask.sum() + green_mask.sum()) / (h * w)
            if nature_ratio > 0.3:
                return ContentType.LANDSCAPE

        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        unique_vals = len(np.unique(gray[::4, ::4]))
        if unique_vals < 50:
            return ContentType.ILLUSTRATION

        return ContentType.PHOTO

    def _estimate_density(self, image: np.ndarray) -> str:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        ratio = np.count_nonzero(edges) / edges.size
        if ratio < 0.03:
            return "sparse"
        elif ratio < 0.08:
            return "normal"
        return "dense"

    def _save_debug(self, name: str, img: np.ndarray, debug_dir: str):
        Path(debug_dir).mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(Path(debug_dir) / f"{name}.png"), img)

    def process(
        self,
        image: np.ndarray,
        config: PipelineConfig | None = None,
        output_path: str = "outputs/result.svg",
    ) -> dict:
        if config is None:
            config = PipelineConfig()

        config.apply_quality()
        config.apply_difficulty()

        if config.save_debug:
            self._save_debug("00_original_input", image, config.debug_dir)

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        if config.content_type == ContentType.AUTO:
            config.content_type = self._detect_content_type(image)
            if config.line_style == LineArtStyle.AUTO:
                style_map = {
                    ContentType.ANIME: LineArtStyle.MANGA,
                    ContentType.ILLUSTRATION: LineArtStyle.INFORMATIVE,
                    ContentType.PORTRAIT: LineArtStyle.FUSED,
                    ContentType.LANDSCAPE: LineArtStyle.FUSED,
                    ContentType.PHOTO: LineArtStyle.FUSED,
                }
                config.line_style = style_map.get(config.content_type, LineArtStyle.FUSED)

        print(f"  [Pipeline] Detected content type: {config.content_type.value}")
        if config.density == "auto":
            config.density = self._estimate_density(image)
        print(f"  [Pipeline] Estimated density: {config.density}")

        if config.k_colors is None:
            config.k_colors = estimate_k(image)
        print(f"  [Pipeline] Target colors: {config.k_colors} (Spatial: {config.spatial_weight})")

        if (
            config.upscale_if_small
            and self.super_res.available
            and self.super_res.should_upscale(image)
        ):
            print("  [Pipeline] Upscaling small image...")
            image = self.super_res.upscale(image)
            if config.save_debug:
                self._save_debug("01_upscaled", image, config.debug_dir)

        h, w = image.shape[:2]
        if max(h, w) > config.max_dimension:
            print(f"  [Pipeline] Resizing to {config.max_dimension}px...")
            scale = config.max_dimension / max(h, w)
            image = cv2.resize(image, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)

        original_for_colors = image.copy()

        fg_mask = None
        if config.remove_background and self.bg_remover.available:
            print("  [Pipeline] Removing background...")
            image, fg_mask = self.bg_remover.remove(image)
            if config.save_debug:
                self._save_debug("02_bg_removed", image, config.debug_dir)

        line_input = image.copy()
        if config.density == "dense":
            line_input = cv2.bilateralFilter(line_input, d=15, sigmaColor=100, sigmaSpace=100)

        print("  [Pipeline] Extracting line art...")
        line_art = self.line_art.extract(
            line_input,
            style=config.line_style,
            density=config.density,
            bridge_gaps=config.bridge_gaps,
        )

        if config.save_debug:
            self._save_debug("03_line_art", line_art, config.debug_dir)

        if config.use_sam and self.sam.available:
            try:
                print(f"  [Pipeline] Segmenting with SAM (Expert Detail)...")
                label_map, palette = self.hybrid_seg.segment(
                    original_for_colors, 
                    k_colors=config.k_colors,
                    spatial_weight=config.spatial_weight
                )
            except Exception as e:
                import traceback
                print(f"  [Pipeline] ✗ SAM failed dramatically:")
                traceback.print_exc()
                print(f"  [Pipeline] Falling back to K-Means quantization...")
                _, palette, label_map = quantize_colors(
                    original_for_colors, 
                    k_colors=config.k_colors,
                    spatial_weight=config.spatial_weight
                )
        else:
            _, palette, label_map = quantize_colors(
                original_for_colors, 
                k_colors=config.k_colors,
                spatial_weight=config.spatial_weight
            )

        if label_map is None:
             _, palette, label_map = quantize_colors(
                 original_for_colors, 
                 k_colors=config.k_colors,
                 spatial_weight=config.spatial_weight
             )

        if config.save_debug:
            debug_q = palette[np.clip(label_map, 0, len(palette) - 1)]
            self._save_debug("04_segmented", debug_q.reshape(image.shape), config.debug_dir)

        if (
            config.detect_faces
            and config.content_type == ContentType.PORTRAIT
            and self.face_parser.available
        ):
            print("  [Pipeline] Parsing facial features...")
            face_map = self.face_parser.parse(original_for_colors)
            if face_map is not None:
                faces = self.face_parser.detect_faces(original_for_colors)
                bbox = faces[0] if faces else None
                label_map = self.face_parser.merge_with_label_map(face_map, label_map, face_bbox=bbox)
                max_lbl = label_map.max() + 1
                if max_lbl > len(palette):
                    extra = np.zeros((max_lbl - len(palette), 3), dtype=np.uint8)
                    for idx in range(len(palette), max_lbl):
                        mask = label_map == idx
                        if np.any(mask):
                            extra[idx - len(palette)] = np.median(original_for_colors[mask], axis=0).astype(np.uint8)
                    palette = np.vstack([palette, extra])
                if config.save_debug:
                    vis_labels = (label_map * 15 % 255).astype(np.uint8)
                    self._save_debug("05_face_parsed", vis_labels, config.debug_dir)

        print("  [Pipeline] Merging line art and regions...")
        merged = merge_line_art_with_regions(label_map, line_art, -1)
        
        # Safety check: if line art wiped out >80% of regions, skip it
        n_orig = len(np.unique(label_map))
        n_merged = len(np.unique(merged[merged >= 0]))
        if n_merged < n_orig * 0.2 and n_orig > 5:
            print(f"  [Pipeline] ⚠ Line art merge lost too much detail ({n_merged}/{n_orig}). Skipping lines for better detail.")
            merged = label_map.copy()

        # Save merged final (after fallback)
        if config.save_debug:
            vis_merged = palette[np.clip(merged, 0, len(palette) - 1)]
            vis_merged = vis_merged.reshape(image.shape)
            vis_merged[merged == -1] = [255, 255, 255]
            self._save_debug("06_merged_final", vis_merged, config.debug_dir)

        # Adaptive min area for small images
        effective_min_area = config.min_region_area
        img_size = max(image.shape[:2])
        if img_size < 512:
            effective_min_area = min(effective_min_area, 50)
            print(f"  [Pipeline] Small image detected, reducing min_area to {effective_min_area}")

        print(f"  [Pipeline] Cleaning regions (min_area={effective_min_area})...")
        cleaned = merge_small_regions_into_neighbors(merged, min_area=effective_min_area, palette=palette)

        # ── FILL GAPS (New) ──
        print("  [Pipeline] Smoothing map and filling internal gaps...")
        cleaned = smooth_label_map(cleaned)

        if fg_mask is not None:
            # Re-apply background mask after smoothing if needed
            cleaned[fg_mask == 0] = -1

        if config.save_debug:
            debug_c = palette[np.clip(cleaned, 0, len(palette) - 1)]
            debug_c = debug_c.reshape(image.shape)
            debug_c[cleaned == -1] = [255, 255, 255]
            self._save_debug("07_cleaned_final", debug_c, config.debug_dir)

        depth_map = None
        difficulty_layers = None
        coloring_order = None
        region_info = None

        if config.use_depth and self.depth.available:
            print("  [Pipeline] Calculating depth and difficulty...")
            try:
                depth_map = self.depth.estimate(original_for_colors)
                difficulty_layers = self.depth.assign_difficulty_layers(depth_map, cleaned, n_layers=config.depth_layers)
                coloring_order = self.depth.create_ordering(depth_map, cleaned)
                region_info = compute_region_difficulty(cleaned, palette, depth_map, config.depth_layers)
                if config.save_debug:
                    depth_vis = (depth_map * 255).astype(np.uint8)
                    depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_VIRIDIS)
                    self._save_debug("06_depth", depth_vis, config.debug_dir)
            except Exception as e:
                print(f"Depth estimation failed: {e}")

        # ── GENERATE SVGS ──
        # 1. Colored SVG
        svg_colored = generate_svg(
            label_map=cleaned,
            palette=palette,
            output_path=output_path,
            min_area=config.min_region_area // 2,
            difficulty_layers=difficulty_layers,
            fill_color=True,
        )

        # 2. Outline SVG
        outline_path = output_path.replace(".svg", "_outline.svg")
        svg_outline = generate_svg(
            label_map=cleaned,
            palette=palette,
            output_path=outline_path,
            min_area=config.min_region_area // 2,
            difficulty_layers=difficulty_layers,
            fill_color=False,
        )

        num_regions = svg_colored.count('class="fillable"')

        # ── ADJACENCY (New) ──
        print("  [Pipeline] Calculating region adjacency...")
        adj_matrix = get_adjacency_matrix(cleaned)

        result = {
            "svg_string": svg_colored,
            "svg_outline_string": svg_outline,
            "svg_path": output_path,
            "svg_outline_path": outline_path,
            "palette": [
                {
                    "index": i + 1,
                    "hex": f"#{int(c[2]):02X}{int(c[1]):02X}{int(c[0]):02X}",
                }
                for i, c in enumerate(palette)
            ],
            "num_regions": int(num_regions),
            "adjacency_matrix": [[int(a), int(b)] for a, b in adj_matrix],
            "k_colors_used": int(len(palette)),
            "density_used": str(config.density),
            "content_type": str(config.content_type.value),
            "models_used": self._get_models_used(config),
        }

        if coloring_order:
            result["suggested_order"] = [int(x) for x in coloring_order]
        if difficulty_layers:
            result["difficulty_layers"] = {int(k): int(v) for k, v in difficulty_layers.items()}
        if region_info:
            result["regions"] = [
                {
                    "label": int(r.label),
                    "area": int(r.area),
                    "color": str(r.color_hex),
                    "difficulty_layer": int(r.difficulty_layer),
                    "suggested_order": int(r.suggested_order),
                    "centroid": [int(r.centroid[0]), int(r.centroid[1])],
                    "seed_point": [int(r.seed_point[0]), int(r.seed_point[1])],
                }
                for r in region_info
            ]

        return result

    def _get_models_used(self, config: PipelineConfig) -> list[str]:
        used = []
        if config.line_style in (LineArtStyle.FUSED, LineArtStyle.INFORMATIVE):
            if self.registry.available("informative_drawings"):
                used.append("informative-drawings")
        if config.line_style in (LineArtStyle.FUSED, LineArtStyle.ANYLINE):
            if self.registry.available("anyline"):
                used.append("anyline")
        if config.line_style == LineArtStyle.MANGA:
            if self.registry.available("manga_line"):
                used.append("manga-line")
        if config.use_sam and self.sam.available:
            used.append("sam2")
        if config.use_depth and self.depth.available:
            used.append("depth-anything-v2")
        if config.content_type == ContentType.PORTRAIT and self.face_parser.available:
            used.append("face-parser")
        return used

    async def generate_from_prompt(
        self,
        prompt: str,
        style: str = "coloring_book",
        config: PipelineConfig | None = None,
        output_path: str = "outputs/generated.svg",
    ) -> dict | None:
        if not self.generator.available:
            return None
        image = await self.generator.generate_from_prompt(prompt, style)
        if image is None:
            return None
        result = self.process(image, config, output_path)
        result["source"] = "generated"
        result["prompt"] = prompt
        return result