"""
AIFinder API Server
Serves classification and training endpoints for the frontend.

Public API:
    POST /v1/classify  — classify text, returns top-N provider predictions.
    No API key required. Rate-limited to 60 requests/minute per IP.
"""

import os
import re
import sys
import json
import joblib
import numpy as np
import torch
import torch.nn as nn
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

from config import MODEL_DIR
from model import AIFinderNet
from features import FeaturePipeline

app = Flask(__name__, static_folder="static", static_url_path="")
CORS(app)
limiter = Limiter(get_remote_address, app=app, default_limits=[])

DEFAULT_TOP_N = 5

pipeline = None
provider_enc = None
net = None
device = None
checkpoint = None


def load_models():
    global pipeline, provider_enc, net, device, checkpoint

    pipeline = joblib.load(os.path.join(MODEL_DIR, "feature_pipeline.joblib"))
    provider_enc = joblib.load(os.path.join(MODEL_DIR, "provider_enc.joblib"))

    checkpoint = torch.load(
        os.path.join(MODEL_DIR, "classifier.pt"),
        map_location="cpu",
        weights_only=True,
    )
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net = AIFinderNet(
        input_dim=checkpoint["input_dim"],
        num_providers=checkpoint["num_providers"],
        hidden_dim=checkpoint["hidden_dim"],
        embed_dim=checkpoint["embed_dim"],
        dropout=checkpoint["dropout"],
    ).to(device)
    net.load_state_dict(checkpoint["state_dict"], strict=False)
    net.eval()


@app.route("/")
def index():
    return send_from_directory("static", "index.html")


@app.route("/api/providers", methods=["GET"])
def get_providers():
    """Return list of available providers."""
    return jsonify({"providers": sorted(provider_enc.classes_.tolist())})


@app.route("/api/classify", methods=["POST"])
def classify():
    """Classify text and return provider predictions."""
    data = request.json
    text = data.get("text", "")

    if len(text) < 20:
        return jsonify({"error": "Text too short (minimum 20 characters)"}), 400

    X = pipeline.transform([text])
    X_t = torch.tensor(X.toarray(), dtype=torch.float32).to(device)

    with torch.no_grad():
        prov_logits = net(X_t)

    prov_proba = torch.softmax(prov_logits.float(), dim=1)[0].cpu().numpy()

    top_prov_idxs = np.argsort(prov_proba)[::-1][:5]
    top_providers = [
        {
            "name": provider_enc.inverse_transform([i])[0],
            "confidence": float(prov_proba[i] * 100),
        }
        for i in top_prov_idxs
    ]

    return jsonify(
        {
            "provider": top_providers[0]["name"],
            "confidence": top_providers[0]["confidence"],
            "top_providers": top_providers,
        }
    )


def _strip_think_tags(text):
    """Remove <think>…</think> (and <thinking>…</thinking>) blocks from input."""
    text = re.sub(r"<think(?:ing)?>.*?</think(?:ing)?>", "", text, flags=re.DOTALL)
    return text.strip()


@app.route("/v1/classify", methods=["POST"])
@limiter.limit("60/minute")
def v1_classify():
    """Public API — classify text and return top-N provider predictions.

    Request JSON:
        text (str):  The text to classify. Any <think>/<thinking> tags will be
                     stripped automatically before classification.
        top_n (int): Number of results to return (default: 5).

    Response JSON:
        provider (str):       Best-matching provider name.
        confidence (float):   Confidence % for the top provider.
        top_providers (list): List of {name, confidence} dicts.

    Rate limit: 60 requests per minute per IP.  No API key required.

    NOTE: If the text you are classifying was produced by a model that emits
          <think> or <thinking> blocks, you should strip those tags BEFORE
          sending the text.  This endpoint does it for you automatically, but
          doing it on your side avoids wasting bytes on the wire.
    """
    data = request.get_json(silent=True)
    if not data or "text" not in data:
        return jsonify({"error": "Request body must be JSON with a 'text' field."}), 400

    raw_text = data["text"]
    text = _strip_think_tags(raw_text)
    top_n = data.get("top_n", DEFAULT_TOP_N)

    if not isinstance(top_n, int) or top_n < 1:
        top_n = DEFAULT_TOP_N

    if len(text) < 20:
        return jsonify({"error": "Text too short (minimum 20 characters after stripping think tags)."}), 400

    X = pipeline.transform([text])
    X_t = torch.tensor(X.toarray(), dtype=torch.float32).to(device)

    with torch.no_grad():
        prov_logits = net(X_t)

    prov_proba = torch.softmax(prov_logits.float(), dim=1)[0].cpu().numpy()

    top_idxs = np.argsort(prov_proba)[::-1][:top_n]
    top_providers = [
        {
            "name": provider_enc.inverse_transform([i])[0],
            "confidence": round(float(prov_proba[i] * 100), 2),
        }
        for i in top_idxs
    ]

    return jsonify(
        {
            "provider": top_providers[0]["name"],
            "confidence": top_providers[0]["confidence"],
            "top_providers": top_providers,
        }
    )


@app.route("/api/correct", methods=["POST"])
def correct():
    """Train on a corrected example."""
    data = request.json
    text = data.get("text", "")
    correct_provider = data.get("correct_provider", "")

    if not text or not correct_provider:
        return jsonify({"error": "Missing text or correct_provider"}), 400

    try:
        prov_idx = provider_enc.transform([correct_provider])[0]
    except ValueError as e:
        return jsonify({"error": f"Unknown provider: {e}"}), 400

    X = pipeline.transform([text])
    X_t = torch.tensor(X.toarray(), dtype=torch.float32).to(device)
    y_prov = torch.tensor([prov_idx], dtype=torch.long).to(device)

    net.train()
    for module in net.modules():
        if isinstance(module, nn.modules.batchnorm._BatchNorm):
            module.eval()

    optimizer = torch.optim.AdamW(net.parameters(), lr=1e-4, weight_decay=1e-4)
    optimizer.zero_grad(set_to_none=True)

    prov_criterion = nn.CrossEntropyLoss()
    prov_logits = net(X_t)
    loss = prov_criterion(prov_logits, y_prov)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
    optimizer.step()

    net.eval()

    checkpoint["state_dict"] = net.state_dict()

    return jsonify({"success": True, "loss": float(loss.item())})


@app.route("/api/save", methods=["POST"])
def save_model():
    """Save the current model state to a file for export."""
    global checkpoint
    data = request.json
    filename = data.get("filename", "aifinder_model.pt")

    save_path = os.path.join(MODEL_DIR, filename)
    torch.save(checkpoint, save_path)

    return jsonify({"success": True, "filename": filename})


@app.route("/models/<filename>")
def download_model(filename):
    """Download exported model file."""
    return send_from_directory(MODEL_DIR, filename)


@app.route("/api/status", methods=["GET"])
def status():
    """Check if models are loaded."""
    return jsonify(
        {
            "loaded": net is not None,
            "device": str(device) if device else None,
        }
    )


if __name__ == "__main__":
    print("Loading models...")
    load_models()
    print(f"Ready on {device}")
    app.run(host="0.0.0.0", port=7860)