from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import torch
import pandas as pd
import numpy as np
import os
import gdown

# -----------------------------
# Download files from Google Drive
# -----------------------------
MODEL_DRIVE_ID = "1qdW9UCJlZSSlSRahx30UrLtZ_3MRMm9e"
CSV_DRIVE_ID = "1lzJU7hjG75uQ4hL02Gzy_vU1FtiHOImT"

if not os.path.exists("trained_model"):
    print("Downloading model from Google Drive...")
    os.makedirs("trained_model", exist_ok=True)
    gdown.download_folder(id=MODEL_DRIVE_ID, output="trained_model", quiet=False)
    print("✓ Model downloaded")

if not os.path.exists("projects.csv"):
    print("Downloading projects.csv from Google Drive...")
    gdown.download(id=CSV_DRIVE_ID, output="projects.csv", quiet=False)
    print("✓ projects.csv downloaded")

# -----------------------------
# App initialization
# -----------------------------
app = FastAPI(
    title="Project Idea Similarity API",
    description="Semantic similarity search for graduation projects",
    version="1.0"
)

def clean_text(text: str) -> str:
    text = str(text).strip()
    text = " ".join(text.split())
    return text

print("Loading model...")
model = SentenceTransformer("trained_model")
print("✓ Model loaded")

print("Loading dataset...")
df = pd.read_csv("projects.csv", encoding="latin1")
df.rename(columns={'Name': 'Name'}, inplace=True)
df["project_id"] = df.index + 1
required_cols = {"Name", "Abstract", "Date"}
print("Number of records:", len(df))
if not required_cols.issubset(df.columns):
    raise RuntimeError("Dataset must contain Name, Abstract, Date columns")

df["Abstract"] = df["Abstract"].apply(clean_text)
projects = df.to_dict(orient="records")

EMBEDDINGS_PATH = "/tmp/embeddings.npy"
try:
    print("Loading embeddings...")
    embeddings = np.load(EMBEDDINGS_PATH)
    embeddings = torch.tensor(embeddings)
    print("✓ Embeddings loaded")
except FileNotFoundError:
    print("Computing embeddings...")
    embeddings = model.encode(
        df["Abstract"].tolist(),
        convert_to_tensor=True,
        show_progress_bar=True
    )
    np.save(EMBEDDINGS_PATH, embeddings.cpu().numpy())
    print("✓ Embeddings computed & saved")

# -----------------------------
# Models
# -----------------------------
class SimilarityRequest(BaseModel):
    abstract: str
    top_k: int = 5

class SimilarProject(BaseModel):
    project_id: int
    project_name: str
    project_abstract: str
    project_date: str
    similarity_score: float

class SimilarityResponse(BaseModel):
    input_abstract: str
    top_k: int
    results: list[SimilarProject]

# -----------------------------
# Routes
# -----------------------------
@app.get("/api/health")
def health():
    return {
        "status": "ok",
        "model_loaded": True,
        "projects_loaded": len(projects)
    }

@app.post("/api/similarity", response_model=SimilarityResponse)
def find_similarity(request: SimilarityRequest):
    if not request.abstract or len(request.abstract.strip()) < 20:
        raise HTTPException(
            status_code=400,
            detail="Project abstract must be at least 20 characters"
        )
    cleaned_input = clean_text(request.abstract)
    input_embedding = model.encode(cleaned_input, convert_to_tensor=True)
    scores = torch.nn.functional.cosine_similarity(
        input_embedding.unsqueeze(0),
        embeddings
    )
    top_k = min(request.top_k, len(projects))
    top_scores, top_indices = torch.topk(scores, k=top_k)
    results = []
    for score, idx in zip(top_scores, top_indices):
        project = projects[idx.item()]
        results.append({
            "project_id": project["project_id"]+1,
            "project_name": project["Name"],
            "project_abstract": project["Abstract"],
            "project_date": str(project["Date"]),
            "similarity_score": round(score.item(), 4)
        })
    return {
        "input_abstract": cleaned_input,
        "top_k": top_k,
        "results": results
    }