main.py

from __future__ import annotations

from contextlib import asynccontextmanager
from typing import Any

from dotenv import load_dotenv

load_dotenv()

from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from sqlalchemy.orm import Session

from azira.auth import (
    create_session,
    get_user_by_email,
    normalize_email,
    password_hash,
    public_user,
    require_user,
    revoke_session,
    utcnow,
    validate_password,
    verify_password,
)
from azira.catalog import get_profile
from azira.database import User, get_db, init_db
from azira.gemini import PestAdvisor
from azira.model import ModelLoadError, PestClassifier, PredictionError, PredictionResult
from azira.pdf_report import build_report_pdf
from azira.settings import (
    CLASSES_PATH,
    CONFIDENCE_THRESHOLD,
    DESIGN_DIR,
    GEMINI_API_KEY,
    GEMINI_MODEL,
    MAX_UPLOAD_BYTES,
    MODEL_DEVICE,
    MODEL_PATH,
    STATIC_DIR,
)


class RegisterPayload(BaseModel):
    email: str
    password: str
    full_name: str = ""


class LoginPayload(BaseModel):
    email: str
    password: str


class ReportPayload(BaseModel):
    language: str = "en"
    prediction: dict[str, Any]
    advice: dict[str, Any]
    image: dict[str, Any] = {}
    model: dict[str, Any] = {}


SUPPORTED_LANGUAGES = {"en", "az", "ru"}


def normalize_ui_language(language: str) -> str:
    return language if language in SUPPORTED_LANGUAGES else "en"


def _candidate_payload(result: PredictionResult) -> list[dict[str, Any]]:
    return [{"label": item.label, "confidence": round(item.confidence, 6)} for item in result.candidates]


def _prediction_payload(result: PredictionResult) -> dict[str, Any]:
    profile = get_profile(result.label)
    return {
        "label": result.label,
        "scientific_name": profile.scientific_name,
        "common_name": profile.common_name,
        "confidence": round(result.confidence, 6),
        "is_confident": result.confidence >= CONFIDENCE_THRESHOLD,
    }


@asynccontextmanager
async def lifespan(app: FastAPI):
    init_db()
    try:
        app.state.classifier = PestClassifier(MODEL_PATH, CLASSES_PATH, MODEL_DEVICE)
        app.state.model_error = None
    except ModelLoadError as exc:
        app.state.classifier = None
        app.state.model_error = str(exc)
    app.state.advisor = PestAdvisor(GEMINI_API_KEY, GEMINI_MODEL)
    yield


app = FastAPI(
    title="AZira AI Pest Manager",
    description="Image-based pest identification and Gemini-assisted field guidance.",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/brand", StaticFiles(directory=DESIGN_DIR), name="brand")


@app.get("/", include_in_schema=False)
async def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/api/health")
async def health(request: Request) -> dict[str, Any]:
    classifier: PestClassifier | None = request.app.state.classifier
    return {
        "ok": classifier is not None,
        "model_loaded": classifier is not None,
        "model_error": request.app.state.model_error,
        "gemini_configured": request.app.state.advisor.is_configured,
        "classes": len(classifier.classes) if classifier else 0,
    }


def auth_response(db: Session, user: User, request: Request) -> dict[str, Any]:
    token = create_session(db, user, request)
    db.refresh(user)
    return {"token": token, "token_type": "bearer", "user": public_user(user)}


@app.post("/api/auth/register")
async def register(request: Request, payload: RegisterPayload, db: Session = Depends(get_db)) -> dict[str, Any]:
    email = normalize_email(payload.email)
    validate_password(payload.password)
    now = utcnow()
    existing = get_user_by_email(db, email)

    if existing:
        raise HTTPException(status_code=409, detail="An account with this email already exists.")

    user = User(
        email=email,
        full_name=payload.full_name.strip()[:160],
        password_hash=password_hash(payload.password),
        created_at=now,
        updated_at=now,
    )
    db.add(user)
    db.commit()
    db.refresh(user)
    return auth_response(db, user, request)


@app.post("/api/auth/login")
async def login(request: Request, payload: LoginPayload, db: Session = Depends(get_db)) -> dict[str, Any]:
    email = normalize_email(payload.email)
    user = get_user_by_email(db, email)
    if not user or not verify_password(payload.password, user.password_hash):
        raise HTTPException(status_code=401, detail="Email or password is incorrect.")
    if not user.is_active:
        raise HTTPException(status_code=403, detail="This account is disabled.")

    return auth_response(db, user, request)


@app.get("/api/auth/me")
async def me(user: User = Depends(require_user)) -> dict[str, Any]:
    return {"user": public_user(user)}


@app.post("/api/auth/logout")
async def logout(authorization: str | None = Header(default=None), db: Session = Depends(get_db)) -> dict[str, bool]:
    if authorization and authorization.lower().startswith("bearer "):
        revoke_session(db, authorization.split(" ", 1)[1].strip())
    return {"ok": True}


@app.get("/api/classes")
async def classes(request: Request) -> dict[str, Any]:
    classifier: PestClassifier | None = request.app.state.classifier
    if classifier is None:
        raise HTTPException(status_code=503, detail=request.app.state.model_error or "Model is not available.")
    return {
        "classes": [
            {
                "label": label,
                "scientific_name": get_profile(label).scientific_name,
                "common_name": get_profile(label).common_name,
            }
            for label in classifier.classes
        ]
    }


@app.post("/api/analyze")
async def analyze(
    request: Request,
    file: UploadFile = File(...),
    crop: str = Form(""),
    region: str = Form(""),
    language: str = Form("en"),
    x_azira_language: str | None = Header(default=None, alias="X-Azira-Language"),
    user: User = Depends(require_user),
) -> dict[str, Any]:
    classifier: PestClassifier | None = request.app.state.classifier
    if classifier is None:
        raise HTTPException(status_code=503, detail=request.app.state.model_error or "Model is not available.")

    content_type = (file.content_type or "").lower()
    if content_type and not content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Upload an image file.")

    image_bytes = await file.read()
    if not image_bytes:
        raise HTTPException(status_code=400, detail="The uploaded image is empty.")
    if len(image_bytes) > MAX_UPLOAD_BYTES:
        raise HTTPException(status_code=413, detail=f"Image is too large. Limit is {MAX_UPLOAD_BYTES // (1024 * 1024)} MB.")

    try:
        result = classifier.predict(image_bytes, top_k=5)
    except PredictionError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    prediction = _prediction_payload(result)
    top_candidates = _candidate_payload(result)
    advisor: PestAdvisor = request.app.state.advisor
    requested_language = (x_azira_language or language or "en").strip()
    language = normalize_ui_language(requested_language)
    print(f"AZira analyze language: form={language!r}, header={x_azira_language!r}, selected={language!r}", flush=True)
    advice = await advisor.advice_for(
        pest_name=result.label,
        confidence=result.confidence,
        crop=crop.strip(),
        region=region.strip(),
        language=language,
        top_candidates=top_candidates,
    )

    return {
        "prediction": prediction,
        "advice": advice,
        "image": {
            "filename": file.filename,
            "content_type": file.content_type,
            "width": result.image_width,
            "height": result.image_height,
        },
        "model": {
            "inference_ms": round(result.inference_ms, 2),
        },
        "language": language,
    }


@app.post("/api/report/pdf")
async def report_pdf(payload: ReportPayload, user: User = Depends(require_user)) -> Response:
    report_data = payload.model_dump()
    report_data["language"] = normalize_ui_language(report_data.get("language", "en"))
    pdf = build_report_pdf(report_data)
    pest_name = report_data.get("prediction", {}).get("label", "pest-report")
    safe_name = "".join(ch if ch.isalnum() else "-" for ch in str(pest_name).lower()).strip("-") or "pest-report"
    return Response(
        content=pdf,
        media_type="application/pdf",
        headers={"Content-Disposition": f'attachment; filename="azira-{safe_name}.pdf"'},
    )