import base64
import io
import urllib.request
from pathlib import Path
from typing import List
from urllib.parse import urlparse

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
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" 

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 get_pytorch_model():
    global _pytorch_model
    if _pytorch_model is None:
        if not PYTORCH_PATH.exists():
            raise FileNotFoundError(f"PyTorch model not found at {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():
            raise FileNotFoundError(f"TensorFlow model not found at {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)  
    arr = arr[:, :, ::-1]         
    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 load_image_from_url(image_url: str):
    if not image_url or not image_url.strip():
        raise ValueError("URL vide ou invalide.")

    if image_url.startswith("data:"):
        try:
            header, encoded = image_url.split(",", 1)
            if "base64" in header:
                image_data = base64.b64decode(encoded)
            else:
                image_data = urllib.request.unquote_to_bytes(encoded)
            return Image.open(io.BytesIO(image_data)).convert("RGB")
        except Exception as exc:
            raise ValueError(f"Impossible de lire le data URL: {exc}")

    parsed = urlparse(image_url)
    if parsed.scheme not in ("http", "https"):
        raise ValueError("L'URL doit commencer par http:// ou https://")

    try:
        request = urllib.request.Request(
            image_url,
            headers={"User-Agent": "GeoClassifier/1.0"},
        )
        with urllib.request.urlopen(request, timeout=15) as response:
            image_data = response.read()
    except Exception as exc:
        raise ValueError(f"Impossible de récupérer l'image depuis l'URL: {exc}")

    try:
        return Image.open(io.BytesIO(image_data)).convert("RGB")
    except Exception as exc:
        raise ValueError(f"Impossible de lire l'image depuis l'URL: {exc}")


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(None),
    image_url: str = Form(None),
    model_choice: str = Form("pytorch")
):
    if image is None and not image_url:
        raise HTTPException(status_code=400, detail="Le fichier ou l'URL est requis.")

    if image is not None:
        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}")
    else:
        try:
            # Accepte URL HTTP/HTTPS, data URL (data:image/...;base64,...) ou base64 pur
            if image_url.startswith("data:"):
                img = load_image_from_url(image_url)
            elif image_url.startswith(("http://", "https://")):
                img = load_image_from_url(image_url)
            else:
                # Assume chaîne base64 pure
                image_data = base64.b64decode(image_url)
                img = Image.open(io.BytesIO(image_data)).convert("RGB")
        except Exception as exc:
            raise HTTPException(status_code=400, detail=f"Impossible de traiter 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)