import io
import os
from pathlib import Path
from typing import List

import numpy as np
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download
import uvicorn

BASE_DIR = Path(__file__).resolve().parent
MODEL_DIR = BASE_DIR / "models"
PYTORCH_PATH = MODEL_DIR / "pytorch_model.pth"
TENSORFLOW_PATH = MODEL_DIR / "model_best.keras"
MODEL_REPO = "danielle2035/intel-classifier-models"
HF_TOKEN = os.environ.get("HF_TOKEN")

CLASSES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
CONFIDENCE_THRESHOLD = 0.6

app = FastAPI(title="Intel Image Classifier")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class CNN(nn.Module):
    def __init__(self, num_classes=6):
        super().__init__()
        self.block1 = self._block(3, 32)
        self.block2 = self._block(32, 64)
        self.block3 = self._block(64, 128)
        self.block4 = self._block(128, 256)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Linear(256, 128)
        self.fc2 = nn.Linear(128, num_classes)
        self.dropout = nn.Dropout(0.5)

    def _block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)
        x = self.gap(x)
        x = x.view(x.size(0), -1)
        x = self.dropout(torch.relu(self.fc1(x)))
        x = self.fc2(x)
        return x


pytorch_transform = transforms.Compose([
    transforms.Resize((150, 150)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

_pytorch_model = None
_tensorflow_model = None


def download_model_file(filename: str, local_path: Path):
    if local_path.exists():
        return
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    try:
        hf_hub_download(
            repo_id=MODEL_REPO,
            filename=filename,
            repo_type="model",
            local_dir=str(MODEL_DIR),
            local_dir_use_symlinks=False,
            use_auth_token=HF_TOKEN,
        )
    except Exception as exc:
        raise FileNotFoundError(
            f"Unable to download {filename} from {MODEL_REPO}: {exc}"
        ) from exc


def get_pytorch_model():
    global _pytorch_model
    if _pytorch_model is None:
        if not PYTORCH_PATH.exists():
            download_model_file("models/pytorch_model.pth", PYTORCH_PATH)
        model = CNN(num_classes=len(CLASSES))
        model.load_state_dict(torch.load(str(PYTORCH_PATH), map_location="cpu"))
        model.eval()
        _pytorch_model = model
    return _pytorch_model


def get_tensorflow_model():
    global _tensorflow_model
    if _tensorflow_model is None:
        if not TENSORFLOW_PATH.exists():
            download_model_file("models/model_best.keras", TENSORFLOW_PATH)
        import tensorflow as tf
        _tensorflow_model = tf.keras.models.load_model(str(TENSORFLOW_PATH), compile=False)
    return _tensorflow_model


def predict_pytorch(image: Image.Image):
    model = get_pytorch_model()
    tensor = pytorch_transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(tensor)
        probs = torch.nn.functional.softmax(outputs, dim=1).squeeze().cpu().numpy()

    sorted_indices = np.argsort(probs)[::-1]
    confidence = float(probs[sorted_indices[0]])
    all_probs = [[CLASSES[i], float(probs[i])] for i in sorted_indices]
    predicted_class = "unknown" if confidence < CONFIDENCE_THRESHOLD else CLASSES[sorted_indices[0]]
    return predicted_class, confidence, all_probs


def predict_tensorflow(image: Image.Image):
    model = get_tensorflow_model()
    img = image.resize((130, 130))
    arr = np.array(img, dtype=np.float32) / 255.0
    arr = np.expand_dims(arr, 0)
    preds = model.predict(arr, verbose=0)[0]
    sorted_indices = np.argsort(preds)[::-1]
    confidence = float(preds[sorted_indices[0]])
    all_probs = [[CLASSES[i], float(preds[i])] for i in sorted_indices]
    predicted_class = "unknown" if confidence < CONFIDENCE_THRESHOLD else CLASSES[sorted_indices[0]]
    return predicted_class, confidence, all_probs


def classify(image: Image.Image, model_choice: str):
    if model_choice == "pytorch":
        return predict_pytorch(image)
    return predict_tensorflow(image)


@app.get("/")
def read_index():
    index_path = BASE_DIR / "index.html"
    if not index_path.exists():
        raise HTTPException(status_code=404, detail="index.html not found")
    return FileResponse(index_path, media_type="text/html")


@app.get("/health")
def health_check():
    return JSONResponse({"status": "ok"})


@app.post("/predict")
def predict(image: UploadFile = File(...), model_choice: str = Form("pytorch")):
    if image.content_type.split('/')[0] != 'image':
        raise HTTPException(status_code=400, detail="Le fichier doit être une image.")

    image_data = image.file.read()
    try:
        img = Image.open(io.BytesIO(image_data)).convert("RGB")
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Impossible de lire l'image: {exc}")

    predicted_class, confidence, all_probs = classify(img, model_choice)
    return {
        "predicted_class": predicted_class,
        "confidence": f"{confidence * 100:.2f}%",
        "probabilities": all_probs,
    }


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)