#import torch and fast api
import torch
from fastapi import FastAPI
from pathlib import Path
from pydantic import BaseModel
from typing import List
import time

#import transformer modules
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

#import sentence_transformers
from sentence_transformers import SentenceTransformer

# import modules
from modules.resume_parser import ResumeParser
from modules.comparison import Comparison

app = FastAPI()

@app.on_event("startup")
def load_models():
    BASE_DIR = Path(__file__).resolve().parent

    # ===== QWEN =====
    base_model_name = "Qwen/Qwen2.5-0.5B-Instruct"
    adapter_path = BASE_DIR / "models/qwen_final"

    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    tokenizer.pad_token = tokenizer.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_name,
        dtype=torch.float32,
        device_map="cpu"
    )

    qwen_model = PeftModel.from_pretrained(base_model, str(adapter_path))
    qwen_model.eval()

    # ===== MiniLM =====
    minilm_model = SentenceTransformer(str(BASE_DIR / "models/mnr_1"))

    # ===== Inject =====
    app.state.resume_parser = ResumeParser(qwen_model, tokenizer)
    app.state.comparison = Comparison(minilm_model)

    print("✅ All models loaded once")


# ---------------------------------------------
# HEALTH CHECK
# ---------------------------------------------
@app.get("/")
def home():
    return {"status": "ATS API Running 🚀"}


# ---------------------------------------------
# MAIN ENDPOINT
# ---------------------------------------------
class ResumeMatchRequest(BaseModel):
    job_description:str
    resume: str

@app.post("/match-resumes/")
async def match_resumes(request:ResumeMatchRequest):
    start = time.time()
    resume = request.resume
    job_description = request.job_description

    #load parser
    parser = app.state.resume_parser
    
    #resumes in parsed format
    resume = parser.format_data(resume)
    
    #resumes after generated
    resume = parser.parse(resume)
    splitted = resume.split("\n")
    name = splitted[0].split("name: ").pop()
    parsed_resume = "\n".join(splitted[1:])
        
    
    #load comparison    
    comparison = app.state.comparison
    
    # use all mini lm to encode
    with torch.no_grad():
        jd = comparison.encode(job_description)
        res = comparison.encode(parsed_resume)
        score = comparison.compare(res, jd)
    
    end = time.time()
    
    return {
        "scores":score,
        "resume":{"name": name, "resume":parsed_resume},
        "threshold":0.670,
        "process_time": f"{end-start} seconds",
        "message": "Candidate Accepted" if score>=0.670 else "Candidate Rejected"
    }