import os, re, json, time, hashlib
from io import BytesIO
from concurrent.futures import ThreadPoolExecutor, as_completed

import pandas as pd
import requests
import streamlit as st

# ─────────────────────────────────────────────
# 1. CONFIGURATION
# ─────────────────────────────────────────────
CATEGORIES = {
    "ADMINISTRATIVE": {"desc": "Notifications, circulars, office orders, memos, meeting minutes, policy instructions",
                       "keywords": ["memo", "circular", "notification", "order", "policy", "minutes", "meeting", "instruction"]},
    "FINANCIAL":      {"desc": "Budget, expenditure, audit, payments, bills, revenue, re-appropriation, PC-1 costing",
                       "keywords": ["budget", "expenditure", "audit", "payment", "revenue", "bill", "fund", "grant", "rupees", "allocation"]},
    "LEGAL":          {"desc": "Contracts, agreements, court cases, writ petitions, judgments, legal opinions",
                       "keywords": ["contract", "agreement", "legal", "court", "judgment", "petition", "writ", "tribunal", "litigation"]},
    "HR":             {"desc": "Recruitment, promotion, transfer/posting, leave, pension, disciplinary matters",
                       "keywords": ["recruitment", "promotion", "transfer", "posting", "leave", "employee", "pension", "appointment", "seniority"]},
    "TECHNICAL":      {"desc": "Tenders, specifications, development projects, IT systems, infrastructure",
                       "keywords": ["specification", "tender", "project", "development", "infrastructure", "procurement", "software", "construction"]},
    "CITIZEN_SERVICE": {"desc": "Citizen applications, complaints, NOCs, requests, grievances",
                       "keywords": ["complaint", "grievance", "noc", "objection", "applicant", "citizen", "request", "redressal"]},
}

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
DEFAULT_MODEL = "openai/gpt-4o-mini"   # check current IDs at openrouter.ai/models

# ─────────────────────────────────────────────
# 2. RULE-BASED ENGINE (optimized v1)
# ─────────────────────────────────────────────
class RuleClassifier:
    """Precompiled, word-boundary regex -> no false hits ('order' in 'border')."""

    def __init__(self, categories):
        self.patterns = {
            cat: re.compile(r"\b(?:" + "|".join(map(re.escape, v["keywords"])) + r")\w*\b", re.IGNORECASE)
            for cat, v in categories.items()
        }

    def classify(self, text):
        hits = {cat: len(p.findall(text)) for cat, p in self.patterns.items()}
        total = sum(hits.values())
        if total == 0:
            return {"category": "UNCLASSIFIED", "confidence": 0.0, "scores": hits}
        best = max(hits, key=hits.get)
        return {"category": best, "confidence": round(hits[best] / total, 2), "scores": hits}


MONTHS = r"(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*"
DATE_RE = re.compile(
    rf"\b(?P<d>\d{{1,2}}(?P<s>[./-])\d{{1,2}}(?P=s)\d{{2,4}}"
    rf"|\d{{1,2}}(?:st|nd|rd|th)?\s+{MONTHS},?\s+\d{{4}}"
    rf"|{MONTHS}\s+\d{{1,2}},?\s+\d{{4}})\b", re.IGNORECASE)
# Reference must contain a digit -> "No objection" is NOT a reference
REF_RE = re.compile(r"\b(?:Ref(?:erence)?|File|Letter|Endst|No)\.?\s*(?:No\.?)?\s*[:#]?\s*"
                    r"([A-Z0-9][\w/&().\-]*\d[\w/&().\-]*)", re.IGNORECASE)

def extract_metadata(text):
    d, r = DATE_RE.search(text), REF_RE.search(text)
    return {"date": d.group("d") if d else None, "reference_no": r.group(1).rstrip(".") if r else None}

# ─────────────────────────────────────────────
# 3. PRIVACY: redact before sending to external AI
# ─────────────────────────────────────────────
PII_PATTERNS = {
    "CNIC":  re.compile(r"\b\d{5}-?\d{7}-?\d\b"),
    "PHONE": re.compile(r"(?:\+92|\b0)3\d{2}[-\s]?\d{7}\b"),
    "EMAIL": re.compile(r"[\w.+-]+@[\w-]+\.[\w.]+"),
}

def redact(text):
    for label, pattern in PII_PATTERNS.items():
        text = pattern.sub(f"[{label}]", text)
    return text

# ─────────────────────────────────────────────
# 4. TEXT EXTRACTION (txt / pdf / docx)
# ─────────────────────────────────────────────
def extract_text(name, data):
    ext = name.lower().rsplit(".", 1)[-1]
    if ext == "txt":
        return data.decode("utf-8", errors="ignore")
    if ext == "pdf":
        from pypdf import PdfReader
        return "\n".join((p.extract_text() or "") for p in PdfReader(BytesIO(data)).pages)
    if ext == "docx":
        from docx import Document
        return "\n".join(p.text for p in Document(BytesIO(data)).paragraphs)
    return ""

# ─────────────────────────────────────────────
# 5. AI ENGINE (OpenRouter)
# ─────────────────────────────────────────────
CAT_LIST = "\n".join(f"- {c}: {v['desc']}" for c, v in CATEGORIES.items())
SYSTEM_PROMPT = f"""You are a records officer AI for the Government of Khyber Pakhtunkhwa, Pakistan.
Classify the document into exactly ONE category:
{CAT_LIST}
Documents may be English, Urdu or mixed, and may contain OCR errors.
Return ONLY valid JSON with these keys:
{{"category": "<one of the categories>", "confidence": <0.0-1.0>, "reason": "<one sentence>",
"subject": "<short subject line>", "summary": "<max 2 sentences>",
"date": "<date or null>", "reference_no": "<file/letter no. or null>",
"department": "<issuing department or null>", "urgency": "Immediate|High|Normal",
"action_required": "<what the receiving officer must do, or null>", "language": "English|Urdu|Mixed"}}"""

def parse_json(content):
    content = re.sub(r"```(?:json)?", "", content).strip()
    start, end = content.find("{"), content.rfind("}")
    return json.loads(content[start:end + 1])

def llm_classify(text, api_key, model, max_chars=6000):
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json",
               "HTTP-Referer": "http://localhost:8501", "X-Title": "KP Document Classifier"}
    payload = {
        "model": model, "temperature": 0,
        "response_format": {"type": "json_object"},
        "messages": [{"role": "system", "content": SYSTEM_PROMPT},
                     {"role": "user", "content": text[:max_chars]}],   # truncate = faster + cheaper
    }
    for attempt in range(3):                                            # retry on rate-limit / server error
        r = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=60)
        if r.status_code == 429 or r.status_code >= 500:
            time.sleep(2 ** attempt)
            continue
        r.raise_for_status()
        data = parse_json(r.json()["choices"][0]["message"]["content"])
        if data.get("category") not in CATEGORIES:
            data["category"] = "UNCLASSIFIED"
        data["confidence"] = max(0.0, min(1.0, float(data.get("confidence", 0))))
        return data
    raise RuntimeError("OpenRouter busy — retried 3 times")

# ─────────────────────────────────────────────
# 6. HYBRID PIPELINE
# ─────────────────────────────────────────────
@st.cache_resource
def get_engine():
    return RuleClassifier(CATEGORIES), {}          # {} = AI response cache (survives reruns)

def classify_one(name, text, cfg, rules, cache):
    t0 = time.perf_counter()
    rule, meta = rules.classify(text), extract_metadata(text)
    res = {"file": name, "words": len(text.split()), "method": "Rule-based",
           "category": rule["category"], "confidence": rule["confidence"],
           "rule_scores": rule["scores"], **meta}

    if not text.strip():
        res["method"] = "Skipped (no text — run OCR Processor first)"
    elif cfg["use_ai"] and cfg["api_key"] and (cfg["always_ai"] or rule["confidence"] < cfg["threshold"]):
        safe_text = redact(text) if cfg["redact"] else text
        key = hashlib.sha256((cfg["model"] + safe_text).encode()).hexdigest()
        try:
            ai = cache.get(key) or llm_classify(safe_text, cfg["api_key"], cfg["model"])
            cache[key] = ai
            res.update({k: v for k, v in ai.items() if v not in (None, "", "null")})
            res["method"] = "AI (OpenRouter)"
        except Exception as e:
            res["method"], res["error"] = "Rule-based (AI failed)", str(e)

    res["latency_s"] = round(time.perf_counter() - t0, 2)
    return res

# ─────────────────────────────────────────────
# 7. STREAMLIT UI
# ─────────────────────────────────────────────
def default_key():
    try:
        return st.secrets["OPENROUTER_API_KEY"]
    except Exception:
        return os.getenv("OPENROUTER_API_KEY", "")

st.set_page_config(page_title="AI Document Classifier v2", page_icon="📄", layout="wide")
st.title("📄 AI Document Classifier v2")
st.caption("Hybrid Rule-based + AI (OpenRouter) • KP Digital Roadmap 2030 — Paperless Government Processes")

with st.sidebar:
    st.header("⚙️ Settings")

    api_key = default_key()                      # loaded silently — never rendered
    if api_key:
        st.success("🔑 API key loaded securely")
    else:
        api_key = st.text_input("OpenRouter API Key", type="password",
                                placeholder="Paste key (not saved)")

    cfg = {
        "api_key":   api_key,
        "model":     st.text_input("Model", value=DEFAULT_MODEL),
        "use_ai":    st.toggle("Use AI", value=True),
        "always_ai": st.toggle("Always use AI (else only when rules unsure)", value=False),
        "threshold": st.slider("Rule confidence needed to skip AI", 0.3, 1.0, 0.7, 0.05),
        "redact":    st.toggle("🔒 Redact CNIC / phone / email before AI", value=True),
        "workers":   st.slider("Parallel workers", 1, 10, 4),
    }
    if cfg["use_ai"] and not cfg["api_key"]:
        st.warning("No API key — running rule-based only.")
        
files = st.file_uploader("Upload documents", type=["txt", "pdf", "docx"], accept_multiple_files=True)

if files and st.button("🚀 Classify", type="primary"):
    rules, cache = get_engine()
    docs = {f.name: extract_text(f.name, f.getvalue()) for f in files}

    results, bar, t0 = [], st.progress(0.0, "Classifying..."), time.perf_counter()
    with ThreadPoolExecutor(max_workers=cfg["workers"]) as pool:
        futures = [pool.submit(classify_one, n, t, cfg, rules, cache) for n, t in docs.items()]
        for i, fut in enumerate(as_completed(futures), 1):
            results.append(fut.result())
            bar.progress(i / len(futures), f"{i}/{len(futures)} done")
    total_time = time.perf_counter() - t0

    df = pd.DataFrame(results)
    c1, c2, c3, c4 = st.columns(4)
    c1.metric("Documents", len(df))
    c2.metric("AI calls", int(df["method"].eq("AI (OpenRouter)").sum()))
    c3.metric("Total time", f"{total_time:.1f}s")
    c4.metric("Avg confidence", f"{df['confidence'].mean():.0%}")

    cols = [c for c in ["file", "category", "confidence", "method", "subject", "urgency",
                        "date", "reference_no", "department", "action_required", "latency_s"] if c in df]
    st.dataframe(df[cols], use_container_width=True, hide_index=True)
    st.bar_chart(df["category"].value_counts())

    low = df[df["confidence"] < 0.6]
    if not low.empty:
        st.warning(f"👤 {len(low)} document(s) below 60% confidence — route for human review.")

    for r in results:
        with st.expander(f"📁 {r['file']} → {r['category']} ({r['confidence']:.0%})"):
            if r.get("summary"): st.info(r["summary"])
            if r.get("reason"):  st.caption(f"Why: {r['reason']}")
            if r.get("error"):   st.error(r["error"])
            st.json({k: v for k, v in r.items() if k not in ("summary", "reason")})
            st.text_area("Text", docs[r["file"]][:3000], height=150, key=f"t_{r['file']}")

    st.download_button("⬇️ Download results (CSV)", df[cols].to_csv(index=False).encode("utf-8"),
                       "classification_results.csv", "text/csv")