import os
import gc

import numpy as np
import pandas as pd
import faiss
import gradio as gr
from fastapi import HTTPException, Query
from fastapi.responses import HTMLResponse

# Your existing helper functions for attribute-based scoring.
from attributes import extract_attributes, attribute_score

# ---------------------------------------------------------------------------
# Paths to pre-built artifacts. Build these ONCE, offline, in a notebook
# (where you have torch/sentence-transformers), then only ship the three
# files below to the Space. Never load a sentence-transformer model here —
# find_company_items only needs precomputed embeddings + a prebuilt FAISS
# index, not a live model.
# ---------------------------------------------------------------------------
DATA_DIR = os.environ.get("DATA_DIR", "data")
MASTER_PATH = os.path.join(DATA_DIR, "master.parquet")
EMB_PATH = os.path.join(DATA_DIR, "embeddings.npy")
INDEX_PATH = os.path.join(DATA_DIR, "faiss_index.bin")

COMPANIES = ["FFL", "PFL", "FFT"]

master: pd.DataFrame | None = None
embeddings = None  # np.memmap, not a full in-RAM copy
index = None  # faiss.Index


def load_resources():
    global master, embeddings, index

    master = pd.read_parquet(
        MASTER_PATH,
        columns=["Item Code", "Company", "Description"],
    )
    master["Item Code"] = master["Item Code"].astype(str).str.strip()
    master["Company"] = master["Company"].astype("category")

    embeddings = np.load(EMB_PATH, mmap_mode="r")
    index = faiss.read_index(INDEX_PATH)

    gc.collect()


# Gradio SDK Spaces have no startup-event hook of their own, so load before
# building/launching the demo, same as the original script did.
load_resources()


def find_company_items(item_code: str, threshold: float = 0.80, top_k: int = 2000):
    source = master[master["Item Code"] == str(item_code).strip()]

    if source.empty:
        return [
            {
                "Company": c,
                "Item Code": "Not Found",
                "Description": "-",
                "Semantic Score": 0,
                "Attribute Score": 0,
                "Final Score": 0,
            }
            for c in COMPANIES
        ]

    idx = source.index[0]
    query_embedding = np.asarray(embeddings[idx]).reshape(1, -1).astype("float32")

    distances, indices = index.search(query_embedding, top_k)

    source_attrs = extract_attributes(source.iloc[0]["Description"])
    results = []

    for company in COMPANIES:
        best = None
        best_score = -1

        for score, i in zip(distances[0], indices[0]):
            if i == -1:
                continue
            row = master.iloc[i]
            if row["Company"] != company:
                continue

            semantic = float(score)
            try:
                attr = attribute_score(source_attrs, extract_attributes(row["Description"]))
            except Exception:
                attr = 1.0

            final = semantic * 0.70 + attr * 0.30

            if final > best_score:
                best_score = final
                best = {
                    "Company": company,
                    "Item Code": row["Item Code"],
                    "Description": row["Description"],
                    "Semantic Score": round(semantic, 4),
                    "Attribute Score": round(attr, 4),
                    "Final Score": round(final, 4),
                }

        if best is None or best_score < threshold:
            results.append(
                {
                    "Company": company,
                    "Item Code": "Not Found",
                    "Description": "-",
                    "Semantic Score": 0,
                    "Attribute Score": 0,
                    "Final Score": 0,
                }
            )
        else:
            results.append(best)

    return sorted(results, key=lambda r: r["Final Score"], reverse=True)


# ---------------------------------------------------------------------------
# Minimal Gradio UI (Gradio SDK Spaces require a `demo` Blocks/Interface to
# launch). This doubles as a simple manual-testing form in the browser.
# ---------------------------------------------------------------------------
def gradio_lookup(item_code, threshold, top_k):
    if master is None:
        return "Server still starting up, try again shortly"
    return find_company_items(item_code, threshold, int(top_k))


demo = gr.Interface(
    fn=gradio_lookup,
    inputs=[
        gr.Textbox(label="Item Code"),
        gr.Slider(0, 1, value=0.80, label="Threshold"),
        gr.Number(value=2000, label="Top K"),
    ],
    outputs=gr.JSON(label="Results"),
    title="Company Item Matcher",
    description="Look up an item code and find matching items across FFL, PFL, and FFT.",
)

# The underlying FastAPI app that Gradio serves. Attaching routes to it is
# the officially supported way to add a plain JSON API alongside the Gradio
# UI on an SDK=gradio Space.
app = demo.app


@app.get("/find-items")
def find_items(
    item_code: str = Query(..., description="Item code to match"),
    threshold: float = 0.80,
    top_k: int = 2000,
):
    if master is None:
        raise HTTPException(status_code=503, detail="Server still starting up, try again shortly")
    return find_company_items(item_code, threshold, top_k)


@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/api", response_class=HTMLResponse)
async def api_info():
    return "<h2>Company Item Matcher API</h2><p>Try /find-items?item_code=YOUR_CODE or /docs</p>"


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)