import io
import copy
import hashlib
import hmac
import json
import os
import random
import re
import shutil
import tempfile
import threading
import uuid
import zipfile
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import torch
from fastapi import FastAPI, File, Form, Header, HTTPException, UploadFile
from huggingface_hub import HfApi
from PIL import Image, ImageOps
from torch import nn
from torch.utils.data import DataLoader, Dataset
from torchvision.models import efficientnet_b0
from torchvision.transforms import v2 as transforms

app = FastAPI(title="Embroidery Complexity API", version="1.0.0")
MODEL_PATH = Path(__file__).parent / "complexidade_best.pth"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load(MODEL_PATH, map_location=device)
classes = checkpoint["classes"]
image_size = checkpoint.get("image_size", 224)


def build_model(model_checkpoint: dict, allow_classifier_mismatch: bool = False) -> torch.nn.Module:
    model_classes = model_checkpoint["classes"]
    trained_model = efficientnet_b0(weights=None)
    trained_model.classifier[1] = torch.nn.Linear(trained_model.classifier[1].in_features, len(model_classes))
    state = model_checkpoint["model_state"]
    if allow_classifier_mismatch:
        current_state = trained_model.state_dict()
        state = {key: value for key, value in state.items() if key in current_state and current_state[key].shape == value.shape}
        trained_model.load_state_dict(state, strict=False)
    else:
        trained_model.load_state_dict(state)
    return trained_model.to(device).eval()


model = build_model(checkpoint)


def make_transform(size: int):
    return transforms.Compose([
    transforms.Resize((size, size)),
    transforms.ToImage(),
    transforms.ToDtype(torch.float32, scale=True),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])


transform = make_transform(image_size)
MODEL_LOCK = threading.RLock()
TRAIN_EXECUTOR = ThreadPoolExecutor(max_workers=1)
TRAINING_JOBS: dict[str, dict] = {}
MAX_DATASET_BYTES = 500 * 1024 * 1024
MAX_EXTRACTED_BYTES = 2 * 1024 * 1024 * 1024
MAX_DATASET_FILES = 5000


def extract_palette(image: Image.Image, quantity: int = 12) -> list[dict]:
    small = image.copy()
    small.thumbnail((512, 512), Image.Resampling.BILINEAR)
    quantized = small.quantize(colors=quantity, method=Image.Quantize.MEDIANCUT)
    palette = quantized.getpalette() or []
    colors = quantized.getcolors(maxcolors=quantity * 512 * 512) or []
    total = sum(count for count, _ in colors) or 1
    result = []
    for count, index in sorted(colors, reverse=True):
        percentage = count / total
        if percentage < 0.01:
            continue
        base = index * 3
        rgb = tuple(int(palette[base + offset]) for offset in range(3))
        result.append({"hex": "#%02x%02x%02x" % rgb, "rgb": list(rgb), "percentage": round(percentage, 4)})
        if len(result) >= 8:
            break
    return result


def _training_token_is_valid(provided_token: str | None) -> bool:
    expected = os.getenv("TRAINING_API_TOKEN", "").strip()
    return bool(expected and provided_token and hmac.compare_digest(provided_token, expected))


def _safe_extract_dataset(zip_bytes: bytes, destination: Path) -> Path:
    """Extract only a bounded ZIP and return the folder containing the labels."""
    if len(zip_bytes) > MAX_DATASET_BYTES:
        raise ValueError("Dataset ZIP is larger than the 500 MB limit.")
    destination.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive:
        entries = [item for item in archive.infolist() if not item.is_dir()]
        if len(entries) > MAX_DATASET_FILES:
            raise ValueError("Dataset contains too many files.")
        total_size = sum(max(0, item.file_size) for item in entries)
        if total_size > MAX_EXTRACTED_BYTES:
            raise ValueError("Uncompressed dataset is larger than the 2 GB limit.")
        destination_root = destination.resolve()
        for item in entries:
            target = (destination / item.filename).resolve()
            if destination_root not in target.parents:
                raise ValueError("Dataset contains an unsafe ZIP path.")
            target.parent.mkdir(parents=True, exist_ok=True)
            with archive.open(item) as source, target.open("wb") as output:
                shutil.copyfileobj(source, output, length=1024 * 1024)

    candidates = [destination, destination / "dataset" / "rotuladas", destination / "rotuladas"]
    for candidate in candidates:
        if all((candidate / label).is_dir() for label in ("baixo_risco", "medio_risco", "alto_risco")):
            return candidate
    raise ValueError("ZIP must contain baixo_risco/, medio_risco/ and alto_risco/ folders.")


def _order_id(path: Path) -> str:
    match = re.match(r"(\d+)_", path.name)
    return match.group(1) if match else path.stem


class TrainingImageDataset(Dataset):
    def __init__(self, samples: list[tuple[Path, int]], image_transform):
        self.samples = samples
        self.transform = image_transform

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        path, label = self.samples[index]
        previous_limit = Image.MAX_IMAGE_PIXELS
        Image.MAX_IMAGE_PIXELS = None
        try:
            with Image.open(path) as source:
                image = ImageOps.exif_transpose(source).convert("RGB")
        finally:
            Image.MAX_IMAGE_PIXELS = previous_limit
        return self.transform(image), label


def _collect_training_samples(root: Path) -> tuple[list[tuple[Path, int]], list[str]]:
    class_names = ["baixo_risco", "medio_risco", "alto_risco"]
    extensions = {".png", ".jpg", ".jpeg", ".webp", ".gif"}
    samples: list[tuple[Path, int]] = []
    for label, class_name in enumerate(class_names):
        folder = root / class_name
        for path in sorted(folder.rglob("*")):
            if path.is_file() and path.suffix.lower() in extensions:
                samples.append((path, label))
    counts = Counter(label for _, label in samples)
    if any(counts[label] < 6 for label in range(len(class_names))):
        raise ValueError("Each class needs at least 6 images.")
    groups: dict[str, list[tuple[Path, int]]] = {}
    for sample in samples:
        groups.setdefault(_order_id(sample[0]), []).append(sample)
    if len(groups) < 9:
        raise ValueError("Dataset needs at least 9 different orders, with 3 per class.")
    for group in groups.values():
        if len({label for _, label in group}) != 1:
            raise ValueError(f"The same order has different labels: {_order_id(group[0][0])}.")
    return samples, class_names


def _split_training_samples(samples: list[tuple[Path, int]]) -> tuple[list, list]:
    groups: dict[str, list[tuple[Path, int]]] = {}
    for sample in samples:
        groups.setdefault(_order_id(sample[0]), []).append(sample)
    by_label: dict[int, list[list[tuple[Path, int]]]] = {}
    for group in groups.values():
        by_label.setdefault(group[0][1], []).append(group)
    rng = random.Random(42)
    training_groups = []
    validation_groups = []
    for label_groups in by_label.values():
        rng.shuffle(label_groups)
        validation_groups.append(label_groups[0])
        training_groups.extend(label_groups[1:])
    flatten = lambda selected: [sample for group in selected for sample in group]
    return flatten(training_groups), flatten(validation_groups)


def _evaluate_training(model_to_test, loader, criterion, class_count: int) -> tuple[float, float, float]:
    model_to_test.eval()
    total_loss = 0.0
    targets: list[int] = []
    predictions: list[int] = []
    with torch.inference_mode():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            logits = model_to_test(images)
            total_loss += criterion(logits, labels).item() * labels.size(0)
            predictions.extend(logits.argmax(1).cpu().tolist())
            targets.extend(labels.cpu().tolist())
    accuracy = sum(actual == predicted for actual, predicted in zip(targets, predictions)) / max(1, len(targets))
    f1_values = []
    for label in range(class_count):
        true_positive = sum(actual == label and predicted == label for actual, predicted in zip(targets, predictions))
        false_positive = sum(actual != label and predicted == label for actual, predicted in zip(targets, predictions))
        false_negative = sum(actual == label and predicted != label for actual, predicted in zip(targets, predictions))
        precision = true_positive / max(1, true_positive + false_positive)
        recall = true_positive / max(1, true_positive + false_negative)
        f1_values.append(2 * precision * recall / max(1e-9, precision + recall))
    return total_loss / max(1, len(loader.dataset)), sum(f1_values) / len(f1_values), accuracy


def _train_and_publish(job_id: str, zip_bytes: bytes, epochs: int, learning_rate: float) -> None:
    global model, checkpoint, classes, image_size, transform
    temporary_root = Path(tempfile.mkdtemp(prefix=f"training_{job_id}_"))
    try:
        TRAINING_JOBS[job_id].update({"status": "running", "message": "Extracting dataset."})
        dataset_root = _safe_extract_dataset(zip_bytes, temporary_root / "dataset")
        samples, class_names = _collect_training_samples(dataset_root)
        training_samples, validation_samples = _split_training_samples(samples)
        if not training_samples or not validation_samples:
            raise ValueError("The dataset split produced an empty training or validation set.")

        weights = checkpoint.get("image_size", 224)
        image_transform = make_transform(int(weights))
        train_loader = DataLoader(TrainingImageDataset(training_samples, image_transform), batch_size=16, shuffle=True, num_workers=0)
        validation_loader = DataLoader(TrainingImageDataset(validation_samples, image_transform), batch_size=16, shuffle=False, num_workers=0)

        candidate = build_model({"classes": class_names, "model_state": checkpoint["model_state"]}, allow_classifier_mismatch=True)
        counts = Counter(label for _, label in training_samples)
        class_weights = torch.tensor(
            [len(training_samples) / (len(class_names) * max(1, counts[label])) for label in range(len(class_names))],
            dtype=torch.float32,
            device=device,
        )
        criterion = nn.CrossEntropyLoss(weight=class_weights)
        optimizer = torch.optim.AdamW(candidate.parameters(), lr=learning_rate, weight_decay=1e-4)
        best_loss = float("inf")
        best_state = None
        no_improvement = 0

        for epoch in range(1, epochs + 1):
            candidate.train()
            running_loss = 0.0
            for images, labels in train_loader:
                images, labels = images.to(device), labels.to(device)
                optimizer.zero_grad()
                loss = criterion(candidate(images), labels)
                loss.backward()
                optimizer.step()
                running_loss += loss.item() * labels.size(0)
            validation_loss, validation_f1, validation_accuracy = _evaluate_training(candidate, validation_loader, criterion, len(class_names))
            TRAINING_JOBS[job_id].update({
                "epoch": epoch,
                "epochs": epochs,
                "train_loss": round(running_loss / max(1, len(train_loader.dataset)), 5),
                "validation_loss": round(validation_loss, 5),
                "validation_f1": round(validation_f1, 5),
                "validation_accuracy": round(validation_accuracy, 5),
            })
            if validation_loss < best_loss - 0.001:
                best_loss = validation_loss
                best_state = copy.deepcopy(candidate.state_dict())
                no_improvement = 0
            else:
                no_improvement += 1
                if no_improvement >= 5:
                    break

        if best_state is None:
            raise RuntimeError("Training did not produce a valid model.")
        candidate.load_state_dict(best_state)
        new_checkpoint = {"model_state": candidate.state_dict(), "classes": class_names, "image_size": int(weights)}
        local_model_path = temporary_root / "complexidade_best.pth"
        torch.save(new_checkpoint, local_model_path)

        hf_token = os.getenv("HF_TOKEN", "").strip() or os.getenv("HUGGINGFACE_HUB_TOKEN", "").strip()
        repo_id = os.getenv("SPACE_ID", "vitordigitizing/embroidery-complexity-api").strip()
        if not hf_token:
            raise RuntimeError("HF_TOKEN Secret is not configured with write permission.")
        TRAINING_JOBS[job_id]["message"] = "Uploading the new model to Hugging Face."
        HfApi(token=hf_token).upload_file(
            path_or_fileobj=str(local_model_path),
            path_in_repo="complexidade_best.pth",
            repo_id=repo_id,
            repo_type="space",
            commit_message=f"Update complexity model from training job {job_id}",
        )
        with MODEL_LOCK:
            checkpoint = new_checkpoint
            classes = class_names
            image_size = int(weights)
            transform = make_transform(image_size)
            model = candidate.to(device).eval()
        TRAINING_JOBS[job_id].update({"status": "completed", "message": "Model trained and uploaded to Hugging Face."})
    except Exception as error:
        TRAINING_JOBS[job_id].update({"status": "failed", "message": str(error)})
    finally:
        shutil.rmtree(temporary_root, ignore_errors=True)


@app.get("/")
def root():
    return {
        "status": "ok",
        "service": "Embroidery Complexity API",
        "model": "embroidery-complexity",
        "endpoints": {
            "health": "/health",
            "predict": "POST /predict with multipart field 'image'",
            "train": "POST /train with a protected ZIP dataset",
            "training_status": "GET /train/{job_id}",
        },
    }


@app.get("/health")
def health():
    return {"status": "ok", "model": "embroidery-complexity", "device": str(device)}


@app.post("/train", status_code=202)
async def train(
    dataset: UploadFile = File(...),
    x_training_token: str | None = Header(default=None),
    epochs: int = Form(default=20),
    learning_rate: float = Form(default=0.0001),
):
    """Start a protected fine-tuning job and publish the best checkpoint."""
    if not _training_token_is_valid(x_training_token):
        raise HTTPException(status_code=401, detail="Invalid training token.")
    if not os.getenv("HF_TOKEN", "").strip() and not os.getenv("HUGGINGFACE_HUB_TOKEN", "").strip():
        raise HTTPException(status_code=503, detail="HF_TOKEN Secret is not configured.")
    if epochs < 1 or epochs > 50:
        raise HTTPException(status_code=400, detail="epochs must be between 1 and 50.")
    if learning_rate <= 0 or learning_rate > 0.01:
        raise HTTPException(status_code=400, detail="learning_rate must be between 0 and 0.01.")
    if not (dataset.filename or "").lower().endswith(".zip"):
        raise HTTPException(status_code=400, detail="Send the dataset as a .zip file.")
    zip_bytes = await dataset.read(MAX_DATASET_BYTES + 1)
    if len(zip_bytes) > MAX_DATASET_BYTES:
        raise HTTPException(status_code=413, detail="Dataset ZIP is larger than the 500 MB limit.")
    job_id = uuid.uuid4().hex
    TRAINING_JOBS[job_id] = {"status": "queued", "message": "Training job queued."}
    TRAIN_EXECUTOR.submit(_train_and_publish, job_id, zip_bytes, epochs, learning_rate)
    return {"status": "queued", "job_id": job_id, "status_url": f"/train/{job_id}"}


@app.get("/train/{job_id}")
def training_status(job_id: str, x_training_token: str | None = Header(default=None)):
    if not _training_token_is_valid(x_training_token):
        raise HTTPException(status_code=401, detail="Invalid training token.")
    job = TRAINING_JOBS.get(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Training job not found.")
    return {"job_id": job_id, **job}


@app.post("/predict")
async def predict(image: UploadFile = File(...)):
    if image.content_type not in {"image/png", "image/jpeg", "image/webp"}:
        raise HTTPException(status_code=400, detail="Only PNG, JPG and WEBP images are supported.")
    try:
        content = await image.read()
        original = ImageOps.exif_transpose(Image.open(io.BytesIO(content))).convert("RGB")
        tensor = transform(original).unsqueeze(0).to(device)
        with torch.inference_mode():
            probabilities = torch.softmax(model(tensor), dim=1)[0]
        index = int(probabilities.argmax())
        return {
            "status": "success",
            "result": classes[index],
            "confidence": round(float(probabilities[index]) * 100, 2),
            "probabilities": {name: round(float(probabilities[pos]) * 100, 2) for pos, name in enumerate(classes)},
            "colors_palette": extract_palette(original),
        }
    except Exception as error:
        raise HTTPException(status_code=400, detail=f"Could not analyze image: {error}") from error