"""
app.py
======
Smart Recruiter — Resume Parser
HuggingFace Space Deployment
Runs with: uvicorn app:app --host 0.0.0.0 --port 7860 --workers 1
"""

import json
import logging
import os
import re
from contextlib import asynccontextmanager

import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from peft import PeftModel
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

ADAPTER_REPO   = "DarkSting/resume_parser"
BASE_MODEL     = "google/gemma-4-E2B-it"
MAX_NEW_TOKENS = 800

SYSTEM_PROMPT = """You are an expert resume parser. Extract structured information from resume text and return ONLY valid JSON with no extra text, no markdown, no explanation.

Always follow this exact schema. Use null for missing fields and [] for missing lists:

{
  "name": "Full Name or null",
  "email": "email@example.com or null",
  "phone": "phone number or null",
  "skills": ["skill1", "skill2"],
  "programming_languages": ["Python", "JavaScript"],
  "experience": [
    {
      "job_title": "Job Title",
      "company": "Company Name",
      "start_date": "MM/YYYY or MM/DD/YYYY or null",
      "end_date": "MM/YYYY or MM/DD/YYYY or Present or null",
      "responsibilities": ["responsibility1", "responsibility2"]
    }
  ],
  "total_experience_years": 0.0,
  "education": [
    {
      "degree": "Degree Name",
      "field": "Field of Study or null",
      "institution": "Institution Name or null",
      "year": 2020
    }
  ],
  "certifications": ["cert1", "cert2"],
  "projects": ["Project Name 1", "Project Name 2"],
  "publications": ["publication1"],
  "achievements": ["achievement1"],
  "summary": "Brief professional summary or null"
}"""

# ---------------------------------------------------------------------------
# Global model state
# ---------------------------------------------------------------------------

model      = None
tokenizer  = None
load_error = None

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

def load_model():
    global model, tokenizer, load_error

    logger.info("Loading base model: %s", BASE_MODEL)
    logger.info("Loading adapter   : %s", ADAPTER_REPO)

    try:
        if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
            torch_dtype = torch.bfloat16
        else:
            torch_dtype = torch.float16

        model_kwargs = dict(
            dtype=torch_dtype,
            device_map="auto",
        )

        model_kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch_dtype,
            bnb_4bit_quant_storage=torch_dtype,
        )

        logger.info("Loading base model ...")
        base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, **model_kwargs)

        logger.info("Loading LoRA adapter ...")
        model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)
        model.eval()

        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
        logger.info("Model ready on device: %s", model.device)

    except Exception as e:
        load_error = str(e)
        logger.exception("Failed to load model: %s", e)


# ---------------------------------------------------------------------------
# Lifespan — load model once at startup
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    load_model()
    yield


# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Smart Recruiter — Resume Parser",
    description="Fine-tuned Gemma-4-E2B-it resume parser. POST raw resume text, get structured JSON.",
    version="1.0.0",
    lifespan=lifespan,
)

# ---------------------------------------------------------------------------
# Schemas
# ---------------------------------------------------------------------------

class AnalyzeRequest(BaseModel):
    content: str


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def parse_resume(resume_text: str) -> dict:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                f"Parse this resume and return ONLY a JSON object:\n\n"
                f"RESUME:\n{resume_text[:3000]}\n\n"
                f"Return ONLY the JSON object, nothing else."
            ),
        },
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(text=text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            temperature=None,
            top_p=None,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated = output[0][inputs["input_ids"].shape[1]:]
    response  = tokenizer.decode(generated, skip_special_tokens=True).strip()

    response = re.sub(r"^```json\s*", "", response)
    response = re.sub(r"```$",        "", response).strip()

    try:
        return json.loads(response)
    except json.JSONDecodeError:
        match = re.search(r"\{[\s\S]*\}", response)
        if match:
            return json.loads(match.group())
        raise ValueError(f"Model output was not valid JSON:\n{response[:300]}")


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/health")
def health():
    return {
        "status":       "ready" if model else "degraded",
        "model_loaded": model is not None,
        "base_model":   BASE_MODEL,
        "adapter_repo": ADAPTER_REPO,
        "load_error":   load_error,
        "device":       str(model.device) if model else None,
    }


@app.post("/analyze")
def analyze(request: AnalyzeRequest):
    if model is None:
        raise HTTPException(
            status_code=503,
            detail=f"Model not loaded. Check /health. Error: {load_error}",
        )

    content = request.content.strip()
    if not content:
        raise HTTPException(
            status_code=400,
            detail="'content' field is empty. Send raw resume text.",
        )

    logger.info("Parsing resume (%d chars) ...", len(content))

    try:
        result = parse_resume(content)
        logger.info("Parse successful. Skills: %s", result.get("skills", []))
        return JSONResponse(content=result)

    except ValueError as e:
        raise HTTPException(status_code=500, detail=str(e))

    except Exception as e:
        logger.exception("Inference error: %s", e)
        raise HTTPException(status_code=500, detail=f"Inference failed: {str(e)}")