# app.py — FastAPI backend with safe model offloading, sliding-window rate limiting, and async history persistence
import asyncio
import json
import logging
import os
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, Dict, List, Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from peft import PeftModel
from pydantic import BaseModel, Field
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# ----- Configuration & Paths -----
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger("colabmind")

BASE_MODEL = "deepseek-ai/deepseek-coder-6.7b-base"
ADAPTER_REPO = "Agasthya0/colabmind-coder-6.7b-ml-qlora"

ROOT = Path(__file__).resolve().parent
STATIC_DIR = ROOT / "static"
STATIC_DIR.mkdir(parents=True, exist_ok=True)

HISTORY_PATH = ROOT / "history.jsonl"
OFFLOAD_DIR = Path("/tmp/llm_offload")
OFFLOAD_DIR.mkdir(parents=True, exist_ok=True)

# Global holder for loaded ML components
ml_resources: Dict[str, Any] = {}

# ----- Rate Limiter (Thread-safe sliding window per IP) -----
RATE_WINDOW = 60        # Seconds
MAX_REQUESTS = 12       # Requests allowed per window
_rate_lock = threading.Lock()
_client_requests: Dict[str, List[float]] = {}


def check_rate_limit(client_ip: str) -> bool:
    now = time.time()
    with _rate_lock:
        timestamps = _client_requests.setdefault(client_ip, [])
        # Drop timestamps outside the active window
        while timestamps and timestamps[0] <= now - RATE_WINDOW:
            timestamps.pop(0)
            
        if len(timestamps) >= MAX_REQUESTS:
            return False
            
        timestamps.append(now)
        return True


# ----- History Persistence -----
_history_lock = threading.Lock()


def append_history(prompt: str, code: str) -> None:
    entry = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "prompt": prompt,
        "code": code,
    }
    try:
        with _history_lock:
            with open(HISTORY_PATH, "a", encoding="utf-8") as f:
                f.write(json.dumps(entry, ensure_ascii=False) + "\n")
    except Exception as e:
        logger.warning("Failed to record query to history: %s", e)


# ----- Model Execution -----
def run_generation(prompt: str, max_new_tokens: int, temperature: float, top_p: float) -> str:
    tokenizer = ml_resources["tokenizer"]
    model = ml_resources["model"]

    if not prompt or not prompt.strip():
        return "⚠️ Empty prompt"

    # Multi-device/CPU/Disk offload handling: target first model layer's device
    target_device = model.device
    if hasattr(model, "hf_device_map") and "" in model.hf_device_map:
        target_device = model.hf_device_map[""]
    elif torch.cuda.is_available():
        target_device = "cuda:0"

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=2048
    ).to(target_device)

    # Hugging Face enforces do_sample=False when temperature == 0.0
    do_sample = temperature > 0.0
    generate_kwargs = {
        "max_new_tokens": max_new_tokens,
        "eos_token_id": tokenizer.eos_token_id,
        "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
        "do_sample": do_sample,
    }

    if do_sample:
        generate_kwargs["temperature"] = temperature
        generate_kwargs["top_p"] = top_p

    with torch.no_grad():
        output_ids = model.generate(**inputs, **generate_kwargs)

    # Strip prompt tokens from output if present
    prompt_len = inputs["input_ids"].shape[-1]
    response_tokens = output_ids[0][prompt_len:]
    output_text = tokenizer.decode(response_tokens, skip_special_tokens=True)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return output_text.strip()


# ----- Lifespan for Clean Startup/Shutdown -----
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Initializing 4-bit BitsAndBytes quantization configuration...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.float16,
        llm_int8_enable_fp32_cpu_offload=True,
    )

    logger.info("Loading tokenizer from %s...", BASE_MODEL)
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    logger.info("Loading base model with auto-dispatch and disk offload...")
    model_base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        offload_folder=str(OFFLOAD_DIR),
        trust_remote_code=True,
    )

    logger.info("Attaching LoRA adapter from %s...", ADAPTER_REPO)
    # Explicitly pass offload_folder to prevent accelerate dispatch crashes on disk-offloaded layers
    model = PeftModel.from_pretrained(
        model_base,
        ADAPTER_REPO,
        device_map="auto",
        offload_folder=str(OFFLOAD_DIR),
    )
    model.eval()

    ml_resources["tokenizer"] = tokenizer
    ml_resources["model"] = model
    logger.info("Model pipeline ready for inference.")

    yield

    ml_resources.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    logger.info("Model unloaded from memory.")


# ----- Application & Routes -----
app = FastAPI(title="ColabMind Coder API", lifespan=lifespan)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


class PredictRequest(BaseModel):
    data: List[Any] = Field(..., min_length=1, description="Expects [prompt, max_tokens, temp, top_p]")


class HistoryItem(BaseModel):
    prompt: Optional[str] = ""
    code: Optional[str] = ""


def _parse_generation_args(data: List[Any]):
    prompt = str(data[0]) if len(data) > 0 and data[0] is not None else ""
    max_tokens = int(data[1]) if len(data) > 1 and data[1] else 512
    temperature = float(data[2]) if len(data) > 2 and data[2] is not None else 0.2
    top_p = float(data[3]) if len(data) > 3 and data[3] is not None else 0.9

    max_tokens = max(16, min(2048, max_tokens))
    temperature = max(0.0, min(2.0, temperature))
    top_p = max(0.0, min(1.0, top_p))

    return prompt, max_tokens, temperature, top_p


@app.get("/", response_class=HTMLResponse)
async def homepage():
    index_file = ROOT / "index.html"
    if index_file.exists():
        return FileResponse(str(index_file))
    return HTMLResponse("<h3>Upload index.html to the application root directory.</h3>", status_code=404)


@app.post("/run/predict")
async def run_predict(req: PredictRequest, request: Request):
    client_ip = request.client.host if request.client else "unknown"
    if not check_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Too many requests — please slow down.")

    try:
        prompt, max_tokens, temperature, top_p = _parse_generation_args(req.data)
        logger.info("Prediction requested | Length: %d | Client: %s", len(prompt), client_ip)

        # Offload heavy synchronous PyTorch call to background thread pool
        output = await asyncio.to_thread(
            run_generation, prompt, max_tokens, temperature, top_p
        )

        # Non-blocking history record
        asyncio.create_task(asyncio.to_thread(append_history, prompt, output))

        return JSONResponse({"data": [output]})
    except Exception as e:
        logger.exception("Inference failed unexpectedly")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/run/explain")
async def run_explain(req: PredictRequest, request: Request):
    client_ip = request.client.host if request.client else "unknown"
    if not check_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Too many requests — please slow down.")

    try:
        raw_code, max_tokens, temperature, top_p = _parse_generation_args(req.data)
        prompt = (
            "### Instruction:\n"
            "Explain the following Python code in detail, line-by-line, and highlight potential bugs or improvements.\n\n"
            f"### Code:\n{raw_code}\n\n"
            "### Explanation:\n"
        )

        output = await asyncio.to_thread(
            run_generation, prompt, max_tokens, temperature, top_p
        )
        return JSONResponse({"data": [output]})
    except Exception as e:
        logger.exception("Code explanation run failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/run/history_append")
async def route_history_append(payload: HistoryItem):
    try:
        await asyncio.to_thread(append_history, payload.prompt or "", payload.code or "")
        return JSONResponse({"ok": True})
    except Exception as e:
        logger.exception("Manual history append failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/run/history_clear")
async def route_history_clear():
    try:
        with _history_lock:
            if HISTORY_PATH.exists():
                HISTORY_PATH.unlink()
        return JSONResponse({"ok": True})
    except Exception as e:
        logger.exception("History clear operation failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health():
    model = ml_resources.get("model")
    if model is None:
        return {"status": "starting", "device": "uninitialized"}

    device_info = getattr(model, "hf_device_map", str(getattr(model, "device", "unknown")))
    return {"status": "ok", "device_distribution": device_info}