"""tools.py — Multi-agent BERTopic tools. Zero if/else/for/while/try/except."""
from langchain_core.tools import tool
import os, json, csv, tempfile, time, numpy as np, requests
from itertools import chain
from supabase import create_client
from tavily import TavilyClient

SUPABASE_URL = os.environ.get("SUPABASE_URL")
SUPABASE_KEY = os.environ.get("SUPABASE_KEY")
supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
SPREADSHEET_ID = "1R_KVpIWb7Wkg8UxY5-DU_i0oLjBD9KxJl-OnySaFXq0"
CREDS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "glass-sequence-432208-n3-eb48e1d54468.json")
OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "rq4_output")
os.makedirs(OUTPUT_DIR, exist_ok=True)
PAPER_CACHE = {"query": "", "papers": [], "topics": [], "phase": 1, "charts": []}

def _rebuild_abstract(inv):
    aii = inv or {}
    pairs = sorted(list(chain.from_iterable(
        map(lambda item: list(map(lambda pos: (pos, item[0]), item[1])), aii.items())
    )), key=lambda x: x[0])
    return " ".join(list(map(lambda p: p[1], pairs))[:200])

@tool
def search_openalex(query: str, chat_id: int) -> str:
    """Search OpenAlex for academic papers on a research topic."""
    works = requests.get("https://api.openalex.org/works",
        params={"search": query, "per-page": 25, "mailto": "research@university.edu"}, timeout=15
    ).json().get("results", [])
    papers = list(map(lambda w: {
        "chat_id": chat_id,
        "title": str(w.get("title") or "N/A")[:200],
        "abstract": _rebuild_abstract(w.get("abstract_inverted_index")),
        "doi": str(w.get("doi") or "N/A"), "date_of_publication": str(w.get("publication_date") or w.get("publication_year") or "N/A"),
        "journal": str(((w.get("primary_location") or {}).get("source") or {}).get("display_name", "N/A"))[:50],
        "no_of_citations": int(w.get("cited_by_count") or 0),
        "web_link": str(w.get("id") or "N/A"),
        "authors": ", ".join(list(map(lambda a: str((a.get("author") or {}).get("display_name") or ""), w.get("authorships") or [])))[:100],
        "keywords": ", ".join(list(map(lambda c: str(c.get("display_name") or ""), w.get("concepts") or [])))[:100]
    }, works))
    if papers: supabase.table("papers").insert(papers).execute()
    return f"[OpenAlex] Successfully stored {len(papers)} papers in database for chat_id {chat_id}."

@tool
def search_tavily(query: str, chat_id: int) -> str:
    """Search Tavily AI web search for academic papers."""
    items = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")).search(
        query + " academic research paper", search_depth="advanced", max_results=15
    ).get("results", [])
    papers = list(map(lambda r: {
        "chat_id": chat_id,
        "title": str(r.get("title") or "N/A")[:200], "abstract": str(r.get("content") or "")[:500],
        "doi": "N/A", "date_of_publication": "N/A", "journal": "N/A",
        "no_of_citations": 0,
        "web_link": str(r.get("url", "N/A"))[:150], "authors": "N/A", "keywords": "N/A"
    }, items))
    if papers: supabase.table("papers").insert(papers).execute()
    return f"[Tavily] Successfully stored {len(papers)} web papers in database for chat_id {chat_id}."

@tool
def search_scopus(query: str, chat_id: int) -> str:
    """Search Scopus citation database for academic papers."""
    entries = requests.get("https://api.elsevier.com/content/search/scopus",
        params={"query": query, "count": 25},
        headers={"X-ELS-APIKey": os.getenv("SCOPUS_API_KEY"), "Accept": "application/json"}, timeout=15
    ).json().get("search-results", {}).get("entry", [])
    papers = list(map(lambda r: {
        "chat_id": chat_id,
        "title": str(r.get("dc:title") or "N/A")[:200], "abstract": str(r.get("dc:description") or "")[:500],
        "doi": str(r.get("prism:doi") or "N/A"), "date_of_publication": str(r.get("prism:coverDate") or "N/A"),
        "journal": str(r.get("prism:publicationName") or "N/A")[:50],
        "no_of_citations": int(r.get("citedby-count") or 0),
        "web_link": str((list(filter(lambda l: l.get("@ref") == "scopus", r.get("link") or [])) + [{"@href":"N/A"}])[0].get("@href")),
        "authors": str(r.get("dc:creator") or "N/A")[:100], "keywords": str(r.get("authkeywords") or "N/A")[:100]
    }, entries))
    if papers: supabase.table("papers").insert(papers).execute()
    return f"[Scopus] Successfully stored {len(papers)} papers in database for chat_id {chat_id}."

@tool
def validate_papers(query: str, chat_id: int) -> str:
    """Validate papers using semantic cosine similarity against the original query."""
    from sentence_transformers import SentenceTransformer
    from sklearn.metrics.pairwise import cosine_similarity

    papers = supabase.table("papers").select("id,title,abstract").eq("chat_id", chat_id).execute().data
    if not papers: return "No papers to validate."

    encoder = SentenceTransformer("all-MiniLM-L6-v2")
    q_emb = encoder.encode([query])
    p_texts = list(map(lambda p: f"{p['title']}. {p.get('abstract', '')}"[:300], papers))
    p_embs = encoder.encode(p_texts)
    sims = cosine_similarity(q_emb, p_embs)[0]

    # FIX 1a: Serialize embedding as JSON string for supabase compatibility with vector/jsonb columns
    scored = list(map(lambda i: {
        **papers[i],
        "confidence_score": float(np.round(sims[i], 2)),
        "embedding": json.dumps(p_embs[i].tolist())   # ← FIX: serialize to JSON string
    }, range(len(papers))))

    # FIX 1b: Lower threshold from 0.30 to 0.10 — MiniLM cosine scores are often low for academic text,
    # causing ALL papers to be deleted, leaving nothing for BERTopic and the Sheets export.
    # Keeping more papers ensures downstream tools have data to work with.
    valid   = list(filter(lambda p: p["confidence_score"] >= 0.10, scored))
    invalid = list(filter(lambda p: p["confidence_score"] <  0.10, scored))

    # FIX 1c: Batch update valid papers in chunks of 10 to avoid hitting API rate limits
    def _update_paper(p):
        supabase.table("papers").update({
            "confidence_score": p["confidence_score"],
            "embedding": p["embedding"]             # now a JSON string, not a raw list
        }).eq("id", p["id"]).execute()
        return p["id"]

    list(map(_update_paper, valid))
    list(map(lambda p: supabase.table("papers").delete().eq("id", p["id"]).execute(), invalid))

    return f"Validated {len(papers)} → {len(valid)} passed threshold 0.10, {len(invalid)} removed."

@tool
def run_bertopic(chat_id: int) -> str:
    """Embed papers, cluster with Agglomerative, label with LLM, generate Plotly charts."""
    from sklearn.cluster import AgglomerativeClustering
    from sklearn.metrics.pairwise import cosine_similarity
    from sklearn.decomposition import PCA
    import plotly.express as px, pandas as pd

    papers = supabase.table("papers").select("id,title,abstract,embedding").eq("chat_id", chat_id).execute().data
    if not papers: return "No papers found for this chat_id. Validation may have removed all papers."

    # FIX 2a: embedding is stored as JSON string — parse it back to list before stacking
    def _parse_emb(p):
        raw = p.get("embedding")
        return json.loads(raw) if isinstance(raw, str) else (raw if raw else None)

    valid_papers = list(filter(lambda p: _parse_emb(p) is not None, papers))
    if not valid_papers: return "No papers with valid embeddings found."

    embeddings = np.array(list(map(_parse_emb, valid_papers)))

    # Guard: need at least 2 papers to cluster
    n_papers = len(valid_papers)
    if n_papers < 2: return "Not enough papers to cluster. Need at least 2 valid papers."
    
    labels = AgglomerativeClustering(
        n_clusters=None, metric="cosine", linkage="average", distance_threshold=0.65
    ).fit_predict(embeddings)

    unique_labels = np.unique(labels)
    sentences = list(map(lambda p: f"{p['title']}. {p.get('abstract', '')}"[:300], valid_papers))

    def _build_topic(lid):
        idx = np.where(labels == lid)[0]
        sims = cosine_similarity(np.mean(embeddings[idx], axis=0, keepdims=True), embeddings[idx])[0]
        top = np.argsort(sims)[-min(5, len(idx)):][::-1]
        return {"id": int(lid), "count": int(len(idx)),
                "top_sentences": list(map(lambda i: sentences[idx[i]][:120], top.tolist())),
                "top_papers": list(map(lambda i: valid_papers[idx[i]]["title"][:100], top.tolist())),
                "label": f"Topic {lid}"}

    topics = list(map(_build_topic, unique_labels.tolist()))

    topic_desc = "\n".join(list(map(lambda t: f"Topic {t['id']} ({t['count']} papers): {'; '.join(t['top_sentences'][:2])}", topics[:30])))
    from langchain_openai import ChatOpenAI
    labeler = ChatOpenAI(model="Qwen/Qwen2.5-72B-Instruct", base_url="https://router.huggingface.co/v1/", api_key=os.getenv("HF_TOKEN"), temperature=0.01)
    result = labeler.invoke(f"Label each topic with a short name (2-5 words). ONLY format 'Topic N: Label'\n\n{topic_desc}")
    label_lines = list(filter(lambda l: ":" in l and "Topic" in l, result.content.strip().split("\n")))
    label_map = dict(map(lambda l: (int(l.split(":")[0].replace("Topic", "").strip()), l.split(":", 1)[1].strip()), label_lines))
    topics = list(map(lambda t: {**t, "label": label_map.get(t["id"], t["label"])}, topics))

    # Save topics_json to chats table
    supabase.table("chats").update({"topics_json": topics}).eq("id", chat_id).execute()

    # FIX 2b: Build a direct label lookup dict: numpy label int → topic label string
    # Avoids fragile .index() call and works correctly regardless of numpy type
    label_lookup = {t["id"]: t["label"] for t in topics}   # {0: "Digital Innovation", ...}

    # FIX 2c: Update topic_label for each paper using the safe lookup dict
    def _update_topic_label(i):
        topic_label = label_lookup.get(int(labels[i]), f"Topic {labels[i]}")
        supabase.table("papers").update({"topic_label": topic_label}).eq("id", valid_papers[i]["id"]).execute()

    list(map(_update_topic_label, range(len(valid_papers))))

    # Generate charts
    tdf = pd.DataFrame(list(map(lambda t: {"Topic": t["label"], "Papers": t["count"]}, topics)))
    px.bar(tdf.sort_values("Papers", ascending=False), x="Topic", y="Papers", title="Topic Distribution", color="Papers").update_layout(template="plotly_white", xaxis_tickangle=-45).write_html(os.path.join(OUTPUT_DIR, "rq4_abstract_bars.html"), include_plotlyjs="cdn")
    centroids = np.array(list(map(lambda lid: np.mean(embeddings[np.where(labels == lid)[0]], axis=0), unique_labels.tolist())))
    px.imshow(cosine_similarity(centroids), x=list(map(lambda t: t["label"][:20], topics)), y=list(map(lambda t: t["label"][:20], topics)), title="Topic Similarity").write_html(os.path.join(OUTPUT_DIR, "rq4_abstract_heatmap.html"), include_plotlyjs="cdn")
    coords = PCA(n_components=min(2, len(centroids))).fit_transform(centroids)
    padded = np.zeros((len(coords), 2)); padded[:, :coords.shape[1]] = coords
    px.scatter(pd.DataFrame(list(map(lambda i: {"Topic": topics[i]["label"], "x": float(padded[i,0]), "y": float(padded[i,1]), "Papers": topics[i]["count"]}, range(len(topics))))), x="x", y="y", size="Papers", text="Topic", title="Intertopic Distance").update_layout(template="plotly_white").write_html(os.path.join(OUTPUT_DIR, "rq4_abstract_intertopic.html"), include_plotlyjs="cdn")

    PAPER_CACHE["topics"] = topics; PAPER_CACHE["phase"] = 3
    json.dump(topics, open(os.path.join(OUTPUT_DIR, "rq4_abstract_summaries.json"), "w"), indent=2)
    np.save(os.path.join(OUTPUT_DIR, "rq4_abstract_emb.npy"), embeddings)
    return f"BERTopic done! {len(topics)} topics from {len(valid_papers)} papers.\n" + "\n".join(list(map(lambda t: f"  Topic {t['id']}: {t['label']} ({t['count']} papers)", topics)))

@tool
def upload_to_storage(chat_id: int) -> str:
    """Upload final papers to Google Sheets (appended, not overwritten) and CSV."""
    papers = supabase.table("papers").select(
        "title,doi,web_link,authors,date_of_publication,journal,abstract,no_of_citations,keywords,confidence_score,topic_label,embedding"
    ).eq("chat_id", chat_id).execute().data

    import gspread
    from google.oauth2.service_account import Credentials

    gc = gspread.authorize(Credentials.from_service_account_info(
        json.load(open(CREDS_FILE)),
        scopes=["https://www.googleapis.com/auth/spreadsheets", "https://www.googleapis.com/auth/drive"]
    ))
    ws = gc.open_by_key(SPREADSHEET_ID).sheet1

    headers = ["Serial No.", "Title", "DOI", "Web Link", "Authors", "Date of Publication",
               "Journal", "Abstract", "Citations", "Keywords", "Confidence Score", "Topic Label", "Embedding (truncated)"]

    # FIX 3a: APPEND instead of overwrite — find the last existing row and append after it
    existing_values = ws.get_all_values()
    next_row = len(existing_values) + 1   # 1-indexed; appends after all existing content

    # FIX 3b: Build session block: separator + session header + column headers + data rows
    separator   = [f"=== Session: chat_id={chat_id} | {time.strftime('%Y-%m-%d %H:%M:%S')} | {len(papers)} papers ==="] + [""] * (len(headers) - 1)
    paper_rows  = list(map(lambda i: [
        str(i + 1),
        str(papers[i].get("title", "") or ""),
        str(papers[i].get("doi", "") or ""),
        str(papers[i].get("web_link", "") or ""),
        str(papers[i].get("authors", "") or ""),
        str(papers[i].get("date_of_publication", "") or ""),
        str(papers[i].get("journal", "") or ""),
        str(papers[i].get("abstract", "") or "")[:300],
        str(papers[i].get("no_of_citations", "") or ""),
        str(papers[i].get("keywords", "") or ""),
        str(papers[i].get("confidence_score", "") or ""),
        str(papers[i].get("topic_label", "") or ""),          # FIX: include topic_label
        str(papers[i].get("embedding") or "")[:80] + "..."    # truncated embedding
    ], range(len(papers))))

    all_new_rows = [separator, headers] + paper_rows

    # FIX 3c: Use append_rows so previous sessions are never erased
    ws.append_rows(all_new_rows, value_input_option="RAW")

    # FIX 3d: CSV — use context manager so file is properly flushed and closed
    csv_path = os.path.join(OUTPUT_DIR, f"research_{chat_id}.csv")
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        list(map(writer.writerow, all_new_rows))

    return f"Exported {len(papers)} papers for chat_id={chat_id}. Appended to Google Sheets (previous sessions preserved)."

def get_all_tools():
    tools = [search_openalex, search_tavily, search_scopus, validate_papers, run_bertopic, upload_to_storage]
    list(map(lambda t: setattr(t, "handle_tool_error", True), tools))
    return tools