src/analyzer/chat/run_chat.py

"""
run_chat.py — unified interactive terminal chatbot for grant data

Features (controlled by flags):
- Startup diagnostics (--verbose)
- LLM-assisted routing (--use-llm-routing, default ON)
- Memory persistence (--with-memory)
- Extended tools (auto-detected or --extended-tools)

Usage:
  python -m src.analyzer.chat.run_chat
  python -m src.analyzer.chat.run_chat --verbose --limit 10
  python -m src.analyzer.chat.run_chat --no-llm-routing
"""
from __future__ import annotations

import argparse
import logging
import os
import sys
import time
from collections import Counter
from pathlib import Path
from typing import List, Dict, Optional
import re

from dotenv import load_dotenv; load_dotenv()

from ..config import load_config
from ..data_loader import load_current_grants, load_past_winners
from ..llm_client import LLMClient
from .chat_tools import ChatTools
from .query_router import route
from ..logging_setup import setup_logging
from ..telemetry.logger import QALogger
from .tool_schemas import openai_tools, detect_extended_features
from ..utils.errors import (
    GrantAnalyzerError,
    ValidationError,
    DataLoadError,
    SearchError,
    LLMError,
    ConfigError
)

# Optional: memory (graceful if not available)
try:
    from .memory import ConversationMemory
    MEMORY_AVAILABLE = True
except ImportError:
    MEMORY_AVAILABLE = False


# ---------------- Domain-term extraction (for dynamic themes) ---------------- #

_STOP = {
    "the","a","an","and","or","of","for","to","in","on","with","by","about","into","from","at","as",
    "call","grant","competition","innovate","uk","round","study","studies","feasibility","phase",
    "funding","programme","program","projects","project","research","development","pilot"
}
_TOKEN_RE = re.compile(r"[a-z0-9][a-z0-9\-]+", re.IGNORECASE)


def _extract_domain_terms(rows: List[Dict], top_n: int = 150) -> set[str]:
    """
    Build a compact vocabulary from your grant corpus (titles/summaries/themes).
    This feeds the router's dynamic theme detection (no hardcoding).
    """
    texts = []
    for r in rows:
        parts = [
            str(r.get("title","")),
            str(r.get("summary","")),
            str(r.get("overview","")),
            str(r.get("scope","")),
            str(r.get("theme","")),
        ]
        texts.append(" ".join(p for p in parts if p))

    unigrams = Counter()
    bigrams  = Counter()

    for txt in texts:
        toks = [t.lower() for t in _TOKEN_RE.findall(txt) if t.lower() not in _STOP and len(t) >= 3]
        unigrams.update(toks)
        for i in range(len(toks)-1):
            w1, w2 = toks[i], toks[i+1]
            if w1 in _STOP or w2 in _STOP:
                continue
            bigrams.update([f"{w1} {w2}"])

    vocab = set([w for w, _ in unigrams.most_common(top_n)])
    vocab |= set([w for w, _ in bigrams.most_common(max(1, top_n // 2))])
    return vocab


def _startup_diagnostics(cfg, *, llm_ok: bool, idx_ok: bool, mem_ok: bool) -> str:
    """Generate startup diagnostics string."""
    provider = getattr(cfg, "provider", "?") if hasattr(cfg, "provider") else cfg.get("provider", "?")
    model = getattr(cfg, "model", "?") if hasattr(cfg, "model") else cfg.get("model", "?")
    openai_key_present = bool(os.getenv("OPENAI_API_KEY"))

    lines = [
        "=== Grant Analyst Chat — Startup Diagnostics ===",
        f"Provider:   {provider}",
        f"Model:      {model}",
        f"LLM Ready:  {llm_ok}",
        f"Index OK:   {idx_ok}",
        f"Memory OK:  {mem_ok}",
        f"API Key:    {'✓' if openai_key_present else '✗'}",
        "",
    ]
    return "\n".join(lines)


def _index_ok_verbose() -> tuple[bool, str]:
    """Check if hybrid index is present and valid."""
    from pathlib import Path
    import pickle

    p = Path("data/index/hybrid_index.pkl")
    if not p.exists():
        return False, f"missing file: {p}"

    try:
        with p.open("rb") as f:
            payload = pickle.load(f)
    except Exception as e:
        return False, f"could not read {p.name}: {e}"

    if not isinstance(payload, dict):
        return False, f"{p.name} is not a dict payload"

    docs = payload.get("docs")
    if not isinstance(docs, list):
        return False, f"{p.name} has no 'docs' list"

    if len(docs) == 0:
        return False, f"{p.name} contains 0 docs"

    return True, f"{p.name} with {len(docs)} docs"


def main(argv: List[str] | None = None) -> None:
    setup_logging()
    cfg = load_config()

    ap = argparse.ArgumentParser(description="Interactive chatbot for grant insights")
    ap.add_argument("--snapshots-dir", type=Path, default=Path("data/snapshots"))
    ap.add_argument("--history-xlsx", type=Path,
                   default=Path("data/IUK-141025-InnovateUKFundedProjects-FY2015-16topresent.xlsx"))
    ap.add_argument("--limit", type=int, default=0,
                   help="Load at most N grants (0 = all)")
    ap.add_argument("--log-jsonl", type=Path, default=Path("_out/chat.jsonl"))

    # Feature flags
    ap.add_argument("--use-llm-routing", dest="use_llm_routing", action="store_true", default=True,
                   help="Enable LLM-assisted routing (default: ON)")
    ap.add_argument("--no-llm-routing", dest="use_llm_routing", action="store_false",
                   help="Disable LLM routing (heuristics only)")
    ap.add_argument("--extended-tools", action="store_true", default=False,
                   help="Force extended tools (insight_search, fetch_link)")
    ap.add_argument("--with-memory", action="store_true", default=False,
                   help="Enable conversation memory (requires memory.py)")
    ap.add_argument("--verbose", action="store_true", default=False,
                   help="Show startup diagnostics")

    args = ap.parse_args(argv)

    # Load data
    logging.info("Loading data ...")
    current = load_current_grants(args.snapshots_dir, limit=args.limit or None)
    past = load_past_winners(args.history_xlsx)
    logging.info("Loaded %d current grants; %d past winners", len(current), len(past))

    # Initialize components
    tools = ChatTools(current, past)

    llm_client = None
    try:
        llm_client = LLMClient(cfg)
    except ConfigError as e:
        logging.warning("LLM config error: %s", e)
    except Exception as e:
        logging.warning("LLMClient init failed: %s", e)

    # JSONL logger
    try:
        args.log_jsonl.parent.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    qalog = QALogger(args.log_jsonl)

    # Index check
    idx_ok, idx_msg = _index_ok_verbose()
    if not idx_ok:
        logging.info("Index check: %s", idx_msg)

    # Memory (optional)
    memory = None
    mem_ok = False
    if args.with_memory and MEMORY_AVAILABLE:
        try:
            memory = ConversationMemory("_out/memory/session.json")
            mem_ok = True
            logging.info("Conversation memory enabled")
        except Exception as e:
            logging.warning("Could not init memory: %s", e)

    # Tool registration
    extended_mode = args.extended_tools or detect_extended_features()
    available_tools = openai_tools(extended=extended_mode)
    logging.info("Registered %d tools (extended: %s)", len(available_tools), extended_mode)

    # Domain terms for routing (if router supports it)
    domain_terms = _extract_domain_terms(current)
    if domain_terms:
        # Try to inject into router (optional feature)
        try:
            from . import query_router
            if hasattr(query_router, 'set_domain_terms'):
                query_router.set_domain_terms(domain_terms)
                query_router.set_fuzzy_threshold(0.84)
                sample = ", ".join(list(sorted(domain_terms, key=len, reverse=True))[:5])
                logging.info("Loaded %d domain terms (e.g., %s ...)", len(domain_terms), sample)
            else:
                logging.debug("Router does not support dynamic domain terms")
        except Exception as e:
            logging.debug("Could not inject domain terms into router: %s", e)

    # Startup diagnostics
    if args.verbose:
        diag = _startup_diagnostics(
            cfg,
            llm_ok=bool(llm_client and llm_client.is_ready()),
            idx_ok=idx_ok,
            mem_ok=mem_ok
        )
        print(diag)

    print("\n💬 Grant Analyst Chat ready! Type 'exit' to quit.\n")

    # REPL loop
    turn_count = 0
    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Goodbye.")
            break

        if not user_input:
            continue

        if user_input.lower() in {"exit", "quit"}:
            print("👋 Goodbye.")
            break

        turn_count += 1
        t0 = time.time()

        # Route intent
        try:
            routed = route(user_input, use_llm=args.use_llm_routing)
        except Exception as e:
            logging.warning("Routing failed (%s). Falling back to heuristic.", e)
            routed = route(user_input, use_llm=False)

        intent = str(routed.get("intent") or "general")
        rargs = routed.get("args") or {}
        answer_md = ""
        ok = True

        try:
            # Handle intents with proper error handling
            if intent in {"search", "list"}:
                filt = (rargs.get("filters") or {}) if isinstance(rargs, dict) else {}
                candidates = rargs.get("keyword_candidates") or []
                primary_kw = rargs.get("keyword") or rargs.get("keyword_hint") or rargs.get("query") or ""
                if primary_kw:
                    candidates = [primary_kw] + [c for c in candidates if c != primary_kw]
                if not candidates:
                    candidates = [""]

                res = []
                used_kw = None
                for kw in candidates:
                    list_kwargs = {
                        "keyword": (kw.strip() if isinstance(kw, str) else ""),
                        "max_award": rargs.get("max_award") or filt.get("max_award"),
                        "audience": rargs.get("audience") or filt.get("audience"),
                        "status": rargs.get("status") or filt.get("status"),  # NEW: Add status filter
                        "limit": rargs.get("limit"),  # FIXED: Don't default to 5, pass None for all
                    }
                    clean_kwargs = {k: v for k, v in list_kwargs.items() if v not in (None, "")}
                    res = tools.list_grants(**clean_kwargs)
                    if res:
                        used_kw = kw
                        break

                if not res:
                    answer_md = "No matching grants found."
                else:
                    hdr = f"### Results (matched on '{used_kw}')" if used_kw else "### Results"
                    # Include status in display
                    bullets = [
                        f"- **{r['id']}** — {r['title']}\n"
                        f"  Status: {r.get('status', 'unknown')} | Deadline: {r.get('deadline','n/a')}"
                        for r in res
                    ]
                    # Add count summary
                    status_counts = {}
                    for r in res:
                        s = r.get('status', 'unknown')
                        status_counts[s] = status_counts.get(s, 0) + 1

                    count_summary = f"\n**Found {len(res)} grant(s)**: " + \
                        ", ".join(f"{count} {status}" for status, count in sorted(status_counts.items()))

                    answer_md = hdr + count_summary + "\n\n" + "\n".join(bullets)

            elif intent == "summarize":
                row = tools.summarize_grant(rargs["grant_id"])
                answer_md = row.get("summary_md", str(row))

            elif intent == "compare":
                diff = tools.compare_grants(rargs["grant_id_a"], rargs["grant_id_b"])
                answer_md = diff.get("comparison_md", str(diff))

            elif intent == "deadlines":
                dl = tools.deadlines_overview(rargs.get("n", 5))
                if not dl:
                    answer_md = "No deadlines available."
                else:
                    answer_md = "### Upcoming deadlines\n" + "\n".join(
                        f"- **{d['title']}** → {d['deadline']}" for d in dl
                    )

            else:
                # General Q&A
                if llm_client and llm_client.is_ready():
                    answer_md = llm_client.summarize(user_input)
                else:
                    answer_md = "LLM not available. Try a structured command like `list battery` or `summarize competition-2316`."

        except ValidationError as e:
            ok = False
            answer_md = f"⚠️  Invalid input: {e}"
            logging.debug("Validation error: %s", e)

        except DataLoadError as e:
            ok = False
            answer_md = f"⚠️  Data error: {e}"
            logging.error("Data load error: %s", e)

        except SearchError as e:
            ok = False
            answer_md = f"⚠️  Search error: {e}"
            logging.error("Search error: %s", e)

        except LLMError as e:
            ok = False
            answer_md = f"⚠️  LLM error: {e}\n💡 Tip: Check your API key and internet connection"
            logging.error("LLM error: %s", e)

        except GrantAnalyzerError as e:
            ok = False
            answer_md = f"⚠️  Error: {e}"
            logging.error("Grant analyzer error: %s", e)

        except Exception as e:
            ok = False
            answer_md = f"❌ Unexpected error: {e}"
            logging.error("Unexpected error in chat turn", exc_info=True)

        latency_ms = int((time.time() - t0) * 1000)
        print(answer_md)

        # Log turn
        try:
            qalog.write(
                user=user_input, intent=intent, args=rargs,
                answer_md=answer_md, ok=ok, latency_ms=latency_ms,
                meta={
                    "model": getattr(llm_client, "model", None),
                    "provider": getattr(llm_client, "provider", None),
                    "use_llm_routing": args.use_llm_routing,
                    "extended_tools": extended_mode
                }
            )
        except Exception:
            pass

        # Update memory (if enabled)
        if memory:
            try:
                memory.add_turn("user", user_input)
                memory.add_turn("assistant", answer_md)

                # Periodic summarization (every 6 turns)
                if turn_count % 6 == 0 and llm_client:
                    memory.update_summary(llm_client.summarize)
            except Exception as e:
                logging.warning("Memory update failed: %s", e)


if __name__ == "__main__":
    main()