from fastapi import FastAPI
from pydantic import BaseModel
from typing import List, Dict
from transformers import AutoTokenizer, AutoModel
import torch

app = FastAPI()

# Load transformer model & tokenizer
MODEL_NAME = "Worko71/keywords_new_model_without_pool" #"Worko71/ketword_match"  # Can be replaced with any HuggingFace model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# Utility: Mean Pooling to get sentence-level embeddings
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]  # First element contains token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
        input_mask_expanded.sum(1), min=1e-9
    )

# ---------- Request Models ----------

class TextPair(BaseModel):
    text1: str
    text2: str

class SkillSet(BaseModel):
    coreSkills: List[str]
    softSkills: List[str]

class JobMap(BaseModel):
    jobs: Dict[str, SkillSet]
    userSkills: SkillSet

class SingleJobComparison(BaseModel):
    userSkills: SkillSet
    jobDetails: SkillSet

# ---------- Endpoint 1: Compare Two Sentences ----------

@app.post("/similarity")
async def compute_similarity(payload: TextPair):
    """
    Compare two texts and return cosine similarity score.
    """
    inputs = tokenizer([payload.text1, payload.text2], padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        model_output = model(**inputs)
    embeddings = mean_pooling(model_output, inputs['attention_mask'])

    sim_score = torch.nn.functional.cosine_similarity(embeddings[0], embeddings[1], dim=0).item()
    return {"similarity_percent": round(sim_score * 100, 2)}

# ---------- Endpoint 2: Compare User vs Job (Soft + Core) ----------

@app.post("/detailed_compare")
async def compare_user_job(payload: SingleJobComparison):
    """
    Compares a user's soft and core skills with a single job's soft and core skills.
    Returns individual similarity scores.
    """
    def encode_average(texts: List[str]) -> torch.Tensor:
        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
        with torch.no_grad():
            model_output = model(**inputs)
        embeddings = mean_pooling(model_output, inputs['attention_mask'])
        return embeddings.mean(dim=0)  # Average across all skill embeddings

    user_core = encode_average(payload.userSkills.coreSkills)
    user_soft = encode_average(payload.userSkills.softSkills)
    job_core = encode_average(payload.jobDetails.coreSkills)
    job_soft = encode_average(payload.jobDetails.softSkills)

    core_score = torch.nn.functional.cosine_similarity(user_core, job_core, dim=0).item()
    soft_score = torch.nn.functional.cosine_similarity(user_soft, job_soft, dim=0).item()

    return {
        "similarity_percent": {
            "core": round(core_score * 100, 2),
            "soft": round(soft_score * 100, 2)
        }
    }

# ---------- Endpoint 3: Bulk Comparison (User vs All Jobs) ----------

@app.post("/bulk_compare")
async def compare_bulk(payload: JobMap):
    """
    Compare a user's skill set against multiple jobs.
    Returns a similarity map per job ID with core and soft skill scores.
    """
    def encode_average(texts: List[str]) -> torch.Tensor:
        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
        with torch.no_grad():
            model_output = model(**inputs)
        embeddings = mean_pooling(model_output, inputs['attention_mask'])
        return embeddings.mean(dim=0)

    user_core = encode_average(payload.userSkills.coreSkills)
    user_soft = encode_average(payload.userSkills.softSkills)

    result = {}
    for job_id, job in payload.jobs.items():
        job_core = encode_average(job.coreSkills)
        job_soft = encode_average(job.softSkills)

        core_score = torch.nn.functional.cosine_similarity(user_core, job_core, dim=0).item()
        soft_score = torch.nn.functional.cosine_similarity(user_soft, job_soft, dim=0).item()

        result[job_id] = {
            "core": round(core_score * 100, 2),
            "soft": round(soft_score * 100, 2)
        }

    return result