import os

import mlflow
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import CamembertForSequenceClassification, CamembertTokenizer

MLFLOW_TRACKING_APP_URI = os.environ["MLFLOW_TRACKING_APP_URI"]

mlflow.set_tracking_uri(MLFLOW_TRACKING_APP_URI)

EXPERIMENT_NAME = "Climate_Fake_News_Detector_Project"
ARTIFACT_PATH = "model"

client = mlflow.tracking.MlflowClient()
experiment = client.get_experiment_by_name(EXPERIMENT_NAME)

if experiment is None:
    raise RuntimeError(f"Experiment MLflow introuvable : {EXPERIMENT_NAME}")

runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["start_time DESC"],
    max_results=1,
)

if len(runs) == 0:
    raise RuntimeError("Aucun run MLflow trouvé pour cet experiment.")

latest_run = runs[0]
run_id = latest_run.info.run_id

local_model_path = mlflow.artifacts.download_artifacts(
    run_id=run_id,
    artifact_path=ARTIFACT_PATH,
)

tokenizer = CamembertTokenizer.from_pretrained(local_model_path)
model = CamembertForSequenceClassification.from_pretrained(local_model_path)
model.eval()

app = FastAPI(
    title="Climate Fake News Detector API",
    description="API de détection de fake news climatiques avec CamemBERT et MLflow.",
    version="1.0.0",
)


class TextInput(BaseModel):
    text: str


@app.get("/")
def home():
    return {
        "message": "Climate Fake News Detector API is running",
        "model_source": "MLflow artifacts",
        "mlflow_run_id": run_id,
    }


@app.post("/predict")
def predict(input_data: TextInput):
    inputs = tokenizer(
        input_data.text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=256,
    )

    with torch.no_grad():
        outputs = model(**inputs)
        prediction = torch.argmax(outputs.logits, dim=1).item()

    label = "fake" if prediction == 1 else "real"

    return {
        "text": input_data.text,
        "prediction": prediction,
        "label": label,
        "mlflow_run_id": run_id,
    }