"""Infer download path templates from existing local model folders."""

from __future__ import annotations

import re
from collections import Counter, defaultdict
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence

from .contracts import ResolvedModel
from .settings import (
    normalize_download_category,
    normalize_relative_subfolder,
    sanitize_folder_name,
)
from .type_utils import normalize_alphanumeric_key

TEMPLATE_CATEGORIES = (
    "loras",
    "checkpoints",
    "embeddings",
    "diffusion_models",
    "text_encoders",
    "controlnet",
    "vae",
    "upscale_models",
)

TAG_FOLDER_TOKENS = {
    "concept",
    "concepts",
    "style",
    "styles",
    "character",
    "characters",
    "clothing",
    "clothes",
    "pose",
    "poses",
    "object",
    "objects",
    "vehicle",
    "vehicles",
    "artist",
    "artists",
    "celebrity",
    "celebrities",
}


def _base_model_token_variants(value: Any) -> List[str]:
    token = normalize_alphanumeric_key(value)
    if not token:
        return []

    variants = [token]
    if token.startswith("flux1") and len(token) > len("flux1"):
        variants.append(f"flux{token[len('flux1'):]}")
    if token.endswith("10") and len(token) > 2:
        variants.append(token[:-1])
    return list(dict.fromkeys(variants))


def _split_relative_folder_segments(relative_path: Any) -> List[str]:
    from .path_utils import split_path_segments
    parts = split_path_segments(relative_path)
    return parts[:-1] if len(parts) > 1 else []


def _get_base_models(base_models_config: Any) -> Iterable[Mapping[str, Any]]:
    if isinstance(base_models_config, Mapping):
        base_models = base_models_config.get("base_models", [])
    else:
        base_models = base_models_config
    if isinstance(base_models, Sequence) and not isinstance(base_models, (str, bytes)):
        for item in base_models:
            if isinstance(item, Mapping):
                yield item


def _build_base_model_alias_index(base_models_config: Any) -> Dict[str, Dict[str, str]]:
    index: Dict[str, Dict[str, str]] = {}
    for item in _get_base_models(base_models_config):
        name = str(item.get("name") or "").strip()
        if not name:
            continue
        aliases = item.get("aliases") or []
        values = [name, *aliases] if isinstance(aliases, list) else [name]
        for value in values:
            for token in _base_model_token_variants(value):
                if token and token not in index:
                    index[token] = {"name": name, "alias": str(value or "")}
    return index


def _match_base_model_segment(
    segment: str,
    base_alias_index: Mapping[str, Dict[str, str]],
    *,
    allow_partial: bool = True,
) -> Optional[Dict[str, str]]:
    token = normalize_alphanumeric_key(segment)
    if not token:
        return None
    if token in base_alias_index:
        return dict(base_alias_index[token])
    if not allow_partial:
        return None

    for alias_token, match in sorted(
        base_alias_index.items(),
        key=lambda item: len(item[0]),
        reverse=True,
    ):
        if len(alias_token) < 4:
            continue
        remainder = token[len(alias_token) :] if token.startswith(alias_token) else ""
        has_version_suffix = bool(re.fullmatch(r"v?\d+[a-z0-9]*", remainder))
        is_plural = token in {f"{alias_token}s", f"{alias_token}es"}
        if has_version_suffix or is_plural:
            return dict(match)
    return None


def _join_mapping_segments(segments: Sequence[str]) -> str:
    return normalize_relative_subfolder("/".join(str(segment or "") for segment in segments))


def _match_base_model_path(
    segments: Sequence[str],
    base_alias_index: Mapping[str, Dict[str, str]],
) -> Optional[Dict[str, Any]]:
    """Find the deepest recognized base-model segment in a folder prefix."""

    best_match: Optional[Dict[str, Any]] = None
    for start_index in range(len(segments)):
        for end_index in range(start_index, len(segments)):
            span = segments[start_index : end_index + 1]
            span_length = end_index - start_index + 1
            match = _match_base_model_segment(
                " ".join(span),
                base_alias_index,
                # A single folder may append a version or plural suffix to an
                # alias (for example IDEOGRAM4 or Upscalers). Multi-folder spans
                # remain exact to avoid guessing across unrelated path parts.
                allow_partial=span_length == 1,
            )
            if not match:
                continue

            current_span_length = int(best_match.get("span_length", 0)) if best_match else 0
            current_index = int(best_match.get("index", -1)) if best_match else -1
            if end_index < current_index:
                continue
            if end_index == current_index and span_length <= current_span_length:
                continue

            best_match = {
                **match,
                "segment": "/".join(span),
                "index": end_index,
                "start_index": start_index,
                "span_length": span_length,
                "path": _join_mapping_segments(segments[: end_index + 1]),
            }
    return best_match


def _is_tag_like_segment(segment: str) -> bool:
    return normalize_alphanumeric_key(segment) in TAG_FOLDER_TOKENS


def _reason_for_template(template: str, confidence: float) -> str:
    percent = round(confidence * 100)
    if template == "{base_model}/{first_tag}":
        return f"{percent}% foldered models use a base-model folder with a second-level grouping."
    if template == "{base_model}":
        return f"{percent}% foldered models use a recognized base-model folder path."
    if template == "{first_tag}":
        return f"{percent}% foldered models start with a tag-like folder."
    if template == "":
        return f"{percent}% models are stored directly in the root folder."
    return f"{percent}% confidence."


def _infer_category_template(records: List[Dict[str, Any]]) -> Dict[str, Any]:
    total = len(records)
    if total == 0:
        return {
            "template": "",
            "confidence": 0.0,
            "apply": False,
            "model_count": 0,
            "foldered_count": 0,
            "reason": "No local models found for this category.",
            "samples": [],
        }

    foldered_count = sum(1 for record in records if record["segments"])
    flat_count = total - foldered_count
    foldered_denominator = max(foldered_count, 1)
    flat_ratio = flat_count / total
    base_first_count = sum(1 for record in records if record["base_model"])
    base_flat_count = sum(
        1
        for record in records
        if record["base_model"] and record.get("segments_after_base", 0) == 0
    )
    base_grouped_count = sum(
        1
        for record in records
        if record["base_model"] and record.get("segments_after_base", 0) >= 1
    )
    tag_first_count = sum(1 for record in records if record["tag_like_first"])

    base_first_ratio = base_first_count / foldered_denominator
    base_flat_ratio = base_flat_count / foldered_denominator
    base_grouped_ratio = base_grouped_count / foldered_denominator
    tag_first_ratio = tag_first_count / foldered_denominator

    template = ""
    confidence = flat_ratio
    if flat_ratio >= 0.65:
        template = ""
        confidence = flat_ratio
    elif base_grouped_ratio >= 0.45 and base_first_ratio >= 0.5:
        template = "{base_model}/{first_tag}"
        confidence = min(1.0, (base_grouped_ratio * 0.75) + (base_first_ratio * 0.25))
    elif base_first_ratio >= 0.45:
        template = "{base_model}"
        confidence = max(base_first_ratio, base_flat_ratio)
    elif tag_first_ratio >= 0.45:
        template = "{first_tag}"
        confidence = tag_first_ratio

    apply = total >= 2 and confidence >= 0.45
    samples = [record["relative_path"] for record in records if record["segments"]][:5]

    return {
        "template": template,
        "confidence": round(confidence, 3),
        "apply": apply,
        "model_count": total,
        "foldered_count": foldered_count,
        "flat_count": flat_count,
        "base_model_count": base_first_count,
        "base_model_folder_count": base_flat_count,
        "base_model_tag_count": base_grouped_count,
        "tag_like_count": tag_first_count,
        "reason": _reason_for_template(template, confidence),
        "samples": samples,
    }


def infer_download_path_templates(
    models: Iterable[ResolvedModel],
    base_models_config: Any,
) -> Dict[str, Any]:
    """Return template and base-model mapping suggestions for local model folders."""

    base_alias_index = _build_base_model_alias_index(base_models_config)
    grouped: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
    mapping_votes: Dict[str, Counter] = defaultdict(Counter)
    seen = set()

    if models is None:
        model_values = ()
    elif isinstance(models, (str, bytes)):
        raise TypeError("models must be an iterable of ResolvedModel values")
    else:
        model_values = models

    for model in model_values:
        if not isinstance(model, ResolvedModel):
            raise TypeError("models must contain ResolvedModel values")
        category = normalize_download_category(str(model.category or ""))
        if category not in TEMPLATE_CATEGORIES:
            continue

        relative_path = str(model.relative_path or model.filename or "")
        model_key = (
            category,
            str(model.path or ""),
            relative_path,
        )
        if model_key in seen:
            continue
        seen.add(model_key)

        segments = _split_relative_folder_segments(relative_path)
        base_match = _match_base_model_path(segments, base_alias_index) if segments else None
        segments_after_base = (
            len(segments) - int(base_match.get("index", 0)) - 1
            if base_match
            else 0
        )
        tag_like_first = _is_tag_like_segment(segments[0]) if segments else False
        grouped[category].append(
            {
                "relative_path": relative_path,
                "segments": segments,
                "base_model": base_match,
                "base_model_path": base_match.get("path", "") if base_match else "",
                "segments_after_base": segments_after_base,
                "tag_like_first": tag_like_first,
            }
        )

        if base_match and segments:
            base_name = base_match.get("name") or ""
            base_path = str(base_match.get("path") or "")
            if base_name and base_path and sanitize_folder_name(base_name, base_name) != base_path:
                mapping_votes[base_name][base_path] += 1

    categories = {
        category: _infer_category_template(grouped.get(category, []))
        for category in TEMPLATE_CATEGORIES
    }
    templates = {
        category: suggestion["template"]
        for category, suggestion in categories.items()
        if suggestion.get("apply")
    }
    base_model_path_mappings = {
        base_name: votes.most_common(1)[0][0]
        for base_name, votes in mapping_votes.items()
        if votes
    }

    return {
        "templates": templates,
        "base_model_path_mappings": base_model_path_mappings,
        "categories": categories,
        "total_models": sum(item.get("model_count", 0) for item in categories.values()),
    }