import os, json, re
from pathlib import Path
import gradio as gr
import numpy as np

# --- Paths
DOCS_DIR = Path("docs")
STATIC_DIR = Path("static")
INDEX_DIR = Path("index"); INDEX_DIR.mkdir(exist_ok=True)
DB_PATH = INDEX_DIR / "faiss.index"
META_PATH = INDEX_DIR / "meta.json"

TITLE = "Rohan Sharma — PersonalGPT (Open-Source)"
INTRO = "Ask about my experience, projects, and fit for your role. Answers cite sources."

# --- Embeddings + FAISS
from sentence_transformers import SentenceTransformer
import faiss

EMBED_NAME = "sentence-transformers/all-MiniLM-L6-v2"
EMBED_DIMS = 384  # MiniLM-L6

def read_text_file(p: Path) -> str:
    return p.read_text(encoding="utf-8", errors="ignore")

def strip_front_matter(text: str):
    if text.startswith("---"):
        parts = text.split("---", 2)
        if len(parts) >= 3:
            return parts[1], parts[2].strip()
    return "", text

def load_corpus():
    blobs = []
    if DOCS_DIR.exists():
        for p in DOCS_DIR.rglob("*.md"):
            text = read_text_file(p)
            _, body = strip_front_matter(text)
            if body.strip():
                blobs.append((str(p), body))
    # (Add PDF extraction later if you want)
    return blobs

def chunk_text(text, chunk_size=800, overlap=150):
    words = text.split()
    chunks, start = [], 0
    while start < len(words):
        end = min(start + chunk_size, len(words))
        chunk = " ".join(words[start:end]).strip()
        if chunk:
            chunks.append(chunk)
        if end == len(words): break
        start = max(0, end - overlap)
    return chunks

EMBED = SentenceTransformer(EMBED_NAME)

def build_or_load_index():
    if DB_PATH.exists() and META_PATH.exists():
        index = faiss.read_index(str(DB_PATH))
        meta = json.loads(META_PATH.read_text())
        return index, meta

    blobs = load_corpus()
    meta, vecs = [], []
    for (fname, text) in blobs:
        for c in chunk_text(text):
            meta.append({"file": fname, "text": c[:1000]})
            v = EMBED.encode([c], normalize_embeddings=True)[0]
            vecs.append(v.astype("float32"))

    if vecs:
        mat = np.vstack(vecs)
        index = faiss.IndexFlatIP(mat.shape[1])
        index.add(mat)
    else:
        index = faiss.IndexFlatIP(EMBED_DIMS)

    faiss.write_index(index, str(DB_PATH))
    META_PATH.write_text(json.dumps(meta, ensure_ascii=False, indent=2))
    return index, meta

INDEX, META = build_or_load_index()

def retrieve(query, k=6):
    if INDEX.ntotal == 0:
        return []
    qv = EMBED.encode([query], normalize_embeddings=True).astype("float32")
    D, I = INDEX.search(qv, k)
    hits = []
    for score, idx in zip(D[0], I[0]):
        if idx == -1: continue
        h = META[idx]
        hits.append({"score": float(score), "file": h["file"], "text": h["text"]})
    return hits

def format_citations(hits):
    uniq = {}
    for h in hits:
        key = h["file"]
        if key not in uniq:
            uniq[key] = h
    lines = []
    for f, h in uniq.items():
        snippet = h["text"][:180].replace("\n", " ")
        lines.append(f"• **{Path(f).name}** — “{snippet}…”")
    return "\n".join(lines) if lines else "_no sources_"

# --- LLM (llama.cpp) on CPU
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

MODEL_CANDIDATES = [
    ("MaziyarPanahi/Phi-3-mini-4k-instruct-gguf", "Phi-3-mini-4k-instruct-Q4_K_M.gguf", "Phi-3-mini-4k-instruct Q4"),
    ("TheBloke/phi-2-GGUF", "phi-2.Q4_K_M.gguf", "Phi-2 Q4"),
    ("unsloth/smollm-1.7b-instruct-GGUF", "smollm-1.7b-instruct.Q4_K_M.gguf", "SmolLM 1.7B Q4"),
]

def load_llm():
    for repo, fname, label in MODEL_CANDIDATES:
        try:
            mp = hf_hub_download(repo_id=repo, filename=fname)
            llm = Llama(
                model_path=mp,
                n_ctx=4096,
                n_threads=max(2, (os.cpu_count() or 4)//2),
                logits_all=False,
                verbose=False
            )
            return llm, label
        except Exception as e:
            print(f"[LLM] failed {repo}/{fname}: {e}")
    raise RuntimeError("No GGUF model could be loaded.")

LLM, LOADED_MODEL = load_llm()

SYSTEM_TMPL = """You are PersonalGPT for {NAME}. Use ONLY the provided context.
Audience: {AUDIENCE}. Tone: {TONE}. Be concise, factual, professional.
If insufficient info, say so and suggest 'Download CV'."""

def make_prompt(user_q, hits, persona, tone):
    ctx = "\n\n".join([f"[Source]\n{h['text']}" for h in hits])
    sys = SYSTEM_TMPL.format(NAME="Rohan Sharma", AUDIENCE=persona, TONE=tone)
    return (
        f"{sys}\n\n"
        f"[User Question]\n{user_q}\n\n"
        f"[Context]\n{ctx}\n\n"
        f"Write:\n- 2–3 sentence executive summary\n- 3–6 bullets with specifics\n"
        f"- Then list 'Citations' (file names + short quotes)."
    )

def llm_generate(prompt):
    out = LLM(
        prompt=prompt,
        max_tokens=512,
        temperature=0.3,
        top_p=0.9,
        repeat_penalty=1.1
    )
    return out["choices"][0]["text"].strip()

def jd_fit(jd_text):
    hits = retrieve(jd_text, k=12)
    score = np.mean([h["score"] for h in hits]) if hits else 0.0
    fit_pct = int(max(0, min(100, (score*100))))  # heuristic
    bag = " ".join([h["text"] for h in hits]).lower()
    words = re.findall(r"[a-zA-Z][a-zA-Z0-9\-\+]{2,}", bag)
    key = sorted(set(words), key=lambda w: bag.count(w), reverse=True)[:10]
    return fit_pct, key, hits

def answer_fn(message, history, persona, tone, show_sources):
    q = message.strip()

    if q.lower().startswith("jd:"):
        jd = q[3:].strip()
        pct, skills, hits = jd_fit(jd)
        citations = format_citations(hits) if show_sources else "_hidden_"
        summary = f"Estimated JD match: **{pct}%**. Alignment on: {', '.join(skills[:6]) or 'n/a'}."
        bullets = "- Heuristic; ask for specifics.\n- See citations/snippets below."
        return f"# Executive Summary\n{summary}\n\n## Details\n{bullets}\n\n## Citations\n{citations}"

    hits = retrieve(q, k=6)
    citations = format_citations(hits) if show_sources else "_hidden_"
    prompt = make_prompt(q, hits, persona, tone)
    text = llm_generate(prompt)

    if "## Citations" not in text:
        text += "\n\n## Citations\n" + citations
    return text

def links_html():
    p = STATIC_DIR / "links.json"
    if p.exists():
        try:
            links = json.loads(p.read_text())
            parts = []
            for label, url in links.items():
                label_disp = "Download CV" if label == "cv" else label.title()
                parts.append(f"<a href='{url}' target='_blank'>{label_disp}</a>")
            return " | ".join(parts)
        except Exception:
            return "_links.json invalid_"
    return "_Add static/links.json to show contact links_"

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# {TITLE}\n{INTRO}")

    with gr.Row():
        with gr.Column(scale=1):
            persona = gr.Radio(
                label="Audience",
                choices=["Hiring Manager", "Recruiter", "CTO/Head of Data"],
                value="Hiring Manager"
            )
            tone = gr.Radio(
                label="Tone",
                choices=["Concise", "Detailed"],
                value="Concise"
            )
            show_sources = gr.Checkbox(label="Show Sources", value=True)
            gr.Markdown("### Suggested prompts\n- Summarize Rohan’s last 5 years\n- Top 5 initiatives with business impact\n- Leadership philosophy in 5 bullets\n- Risks handled and mitigations\n- Fit for Head of Data role?\n- Prefix a JD with `JD:` to get a fit score")

            gr.HTML("<hr/>")
            gr.Markdown(links_html())

        with gr.Column(scale=3):
            chat = gr.ChatInterface(
                fn=answer_fn,
                additional_inputs=[persona, tone, show_sources],
                type="markdown",
                fill_height=True,
                autofocus=True,
                analytics_enabled=False,
                clear_btn="Clear"
            )

    # Footer with diagnostics
    total_chunks = META_PATH.exists() and len(json.loads(META_PATH.read_text())) or 0
    updated = "unknown"
    rm = DOCS_DIR / "resume.md"
    if rm.exists():
        try:
            fm, _ = strip_front_matter(read_text_file(rm))
            m = re.search(r"updated:\s*\"?([\d\-]+)\"?", fm or "")
            updated = m.group(1) if m else "unknown"
        except Exception:
            pass

    gr.Markdown(f"**Model:** {LOADED_MODEL} · **Chunks:** {total_chunks} · **Content updated:** {updated}")

if __name__ == "__main__":
    demo.launch()