import os, io, base64, time, json, re
from pathlib import Path
from typing import List, Optional
import numpy as np

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from groq import Groq

# ── Optional heavy deps (graceful fallback) ───────────────────────
try:
    import fitz  # PyMuPDF
    HAS_FITZ = True
except ImportError:
    HAS_FITZ = False

try:
    from sentence_transformers import SentenceTransformer
    import faiss
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    HAS_EMBEDDER = True
except ImportError:
    HAS_EMBEDDER = False

try:
    import cv2
    HAS_CV2 = True
except ImportError:
    HAS_CV2 = False

from PIL import Image

# ── In-memory knowledge base ──────────────────────────────────────
KB = {
    "files": {},       # filename -> {type, summary, entities, chunks, thumb_b64}
    "chunks": [],      # [{text, source, page, embedding}]
    "index": None,     # FAISS index
}

CHUNK_SIZE = 600  # words per chunk
CHUNK_OVERLAP = 80

# ── Helpers ───────────────────────────────────────────────────────
def chunk_text(text: str, source: str, page: int = 0) -> List[dict]:
    words = text.split()
    chunks = []
    for i in range(0, len(words), CHUNK_SIZE - CHUNK_OVERLAP):
        chunk = " ".join(words[i:i + CHUNK_SIZE])
        if chunk.strip():
            chunks.append({"text": chunk, "source": source, "page": page})
    return chunks

def embed_chunks(chunks: List[dict]) -> np.ndarray:
    if not HAS_EMBEDDER:
        return np.random.rand(len(chunks), 384).astype("float32")
    texts = [c["text"] for c in chunks]
    return embedder.encode(texts, convert_to_numpy=True).astype("float32")

def rebuild_index():
    if not KB["chunks"]:
        KB["index"] = None
        return
    vecs = np.array([c["embedding"] for c in KB["chunks"]], dtype="float32")
    dim = vecs.shape[1]
    if HAS_EMBEDDER:
        index = faiss.IndexFlatIP(dim)
        faiss.normalize_L2(vecs)
        index.add(vecs)
        KB["index"] = index
    else:
        KB["index"] = None

def retrieve(query: str, top_k: int = 8) -> List[dict]:
    if not KB["chunks"]:
        return []
    
    # Vector search if available
    vector_results = []
    if KB["index"] is not None and HAS_EMBEDDER:
        q_vec = embedder.encode([query], convert_to_numpy=True).astype("float32")
        faiss.normalize_L2(q_vec)
        _, indices = KB["index"].search(q_vec, min(top_k, len(KB["chunks"])))
        vector_results = [KB["chunks"][i] for i in indices[0] if i < len(KB["chunks"])]

    # Keyword fallback — always run to catch what vector search misses
    query_words = set(query.lower().split())
    keyword_results = []
    for chunk in KB["chunks"]:
        chunk_lower = chunk["text"].lower()
        # score by how many query words appear in chunk
        score = sum(1 for w in query_words if len(w) > 3 and w in chunk_lower)
        if score > 0:
            keyword_results.append((score, chunk))
    keyword_results.sort(key=lambda x: x[0], reverse=True)
    keyword_chunks = [c for _, c in keyword_results[:top_k]]

    # Merge — deduplicate by text, prioritize vector results
    seen = set()
    merged = []
    for chunk in vector_results + keyword_chunks:
        key = chunk["text"][:100]
        if key not in seen:
            seen.add(key)
            merged.append(chunk)
    
    return merged[:top_k + 4]  # return a few extra for better coverage

def pil_to_b64(pil_img, max_size=1024) -> str:
    img = pil_img.copy()
    img.thumbnail((max_size, max_size), Image.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return base64.b64encode(buf.getvalue()).decode()

def claude_client(api_key: str):
    return Groq(api_key=api_key)

def claude_vision(client, b64_img: str, prompt: str) -> str:
    msg = client.chat.completions.create(
        model="meta-llama/llama-4-scout-17b-16e-instruct",
        max_tokens=1000,
        messages=[{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_img}"}},
            {"type": "text", "text": prompt}
        ]}]
    )
    return msg.choices[0].message.content

def claude_text(client, prompt: str) -> str:
    msg = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        max_tokens=1500,
        messages=[{"role": "user", "content": prompt}]
    )
    return msg.choices[0].message.content

# ── File processors ───────────────────────────────────────────────
def process_pdf(raw: bytes, filename: str, client) -> dict:
    chunks = []
    thumb_b64 = None
    if not HAS_FITZ:
        return {"chunks": [{"text": "PDF processing unavailable (PyMuPDF not installed)", "source": filename, "page": 0}], "thumb_b64": None, "raw_text": ""}
    doc = fitz.open(stream=raw, filetype="pdf")
    full_text = ""
    for page_num, page in enumerate(doc):
        # extract text with better layout preservation
        text = page.get_text("text")
        if not text.strip():
            # try blocks mode for scanned/structured PDFs
            text = page.get_text("blocks")
            if isinstance(text, list):
                text = "\n".join([b[4] for b in text if len(b) > 4 and isinstance(b[4], str)])
        full_text += f"\n[Page {page_num + 1}]\n{text}"
        page_chunks = chunk_text(text, filename, page_num + 1)
        chunks.extend(page_chunks)
        if page_num == 0:
            pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            thumb_b64 = pil_to_b64(img, 300)
    # store up to 80000 chars for large multi-section documents
    return {"chunks": chunks, "thumb_b64": thumb_b64, "raw_text": full_text[:80000]}

def process_image(raw: bytes, filename: str, client) -> dict:
    pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
    b64 = pil_to_b64(pil_img)
    thumb_b64 = pil_to_b64(pil_img, 300)
    description = claude_vision(client, b64, "Describe this image in detail. Extract any text visible. Note objects, people, scenes, data, charts, or diagrams present.")
    chunks = chunk_text(description, filename, 0)
    return {"chunks": chunks, "thumb_b64": thumb_b64, "raw_text": description[:5000]}

def process_video(raw: bytes, filename: str, client) -> dict:
    if not HAS_CV2:
        return {"chunks": [{"text": "Video processing unavailable", "source": filename, "page": 0}], "thumb_b64": None, "raw_text": ""}
    import tempfile
    suffix = Path(filename).suffix or ".mp4"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(raw)
        tmp_path = tmp.name
    try:
        cap = cv2.VideoCapture(tmp_path)
        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        sample_at = np.linspace(0, max(total - 1, 0), min(6, total), dtype=int)
        descriptions = []
        thumb_b64 = None
        for i, idx in enumerate(sample_at):
            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
            ret, frame = cap.read()
            if not ret:
                continue
            pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            b64 = pil_to_b64(pil_img)
            if i == 0:
                thumb_b64 = pil_to_b64(pil_img, 300)
            desc = claude_vision(client, b64, f"Frame {i+1} of a video. Describe what's happening. Note any text, people, objects, or key events.")
            descriptions.append(f"[Frame {i+1}] {desc}")
        cap.release()
        os.unlink(tmp_path)
        full_text = "\n\n".join(descriptions)
        chunks = chunk_text(full_text, filename, 0)
        return {"chunks": chunks, "thumb_b64": thumb_b64, "raw_text": full_text[:5000]}
    except Exception as e:
        os.unlink(tmp_path)
        raise e

def process_text(raw: bytes, filename: str) -> dict:
    text = raw.decode("utf-8", errors="ignore")
    chunks = chunk_text(text, filename, 0)
    return {"chunks": chunks, "thumb_b64": None, "raw_text": text[:5000]}

def summarize_section(client, section_text: str, section_label: str) -> str:
    """Summarize a single section/batch of text."""
    prompt = f"""Summarize this section of a document ({section_label}) in 3-4 sentences. 
Capture all distinct topics, subjects, or subsections mentioned. Be specific, not generic.

Text:
{section_text}"""
    try:
        return claude_text(client, prompt)
    except Exception as e:
        return f"[Could not summarize this section: {e}]"

def generate_summary_and_entities(client, raw_text: str, filename: str) -> dict:
    BATCH_SIZE = 15000  # chars per batch, safely under token limits

    if len(raw_text) <= BATCH_SIZE:
        # Short document — single pass
        content = raw_text
        combined_summary_input = content
    else:
        # Long document — map-reduce: summarize each batch, then combine
        batches = [raw_text[i:i + BATCH_SIZE] for i in range(0, len(raw_text), BATCH_SIZE)]
        section_summaries = []
        for idx, batch in enumerate(batches):
            label = f"part {idx + 1} of {len(batches)}"
            s = summarize_section(client, batch, label)
            section_summaries.append(f"[{label}]: {s}")
        combined_summary_input = "\n\n".join(section_summaries)

    prompt = f"""Based on the following content (or section summaries) from "{filename}", produce a complete analysis.
This document may cover MULTIPLE topics/sections/subjects — make sure your output covers ALL of them, not just the first part.

Respond in JSON only (no markdown):
{{
  "summary": "Detailed 5-8 sentence summary that covers ALL major sections/subjects/topics found in the document",
  "key_entities": ["entity1", "entity2", ...up to 20 entities, covering the whole document],
  "topics": ["topic1", "topic2", ...up to 10 topics covering the full scope],
  "file_type_detected": "what kind of document this is",
  "key_sections": ["list ALL section/chapter/subject titles found in the document"],
  "important_facts": ["fact1", "fact2", ...up to 10 specific facts, drawn from across the ENTIRE document"]
}}

Content:
{combined_summary_input[:14000]}"""
    try:
        result = claude_text(client, prompt)
        clean = result.strip()
        if clean.startswith("```"):
            clean = re.sub(r"```(?:json)?", "", clean).strip().rstrip("`").strip()
        return json.loads(clean)
    except Exception:
        return {"summary": "Could not generate summary.", "key_entities": [], "topics": [], "file_type_detected": "unknown"}

def find_connections(client) -> List[dict]:
    if len(KB["files"]) < 2:
        return []
    summaries = "\n".join([f"- {name}: {info.get('summary','')}" for name, info in KB["files"].items()])
    prompt = f"""Given these documents in a knowledge base, find meaningful connections between them.
Respond in JSON only (no markdown):
[{{"doc1": "filename1", "doc2": "filename2", "connection": "brief description of how they relate"}}]

Documents:
{summaries}"""
    try:
        result = claude_text(client, prompt)
        clean = result.strip()
        if clean.startswith("```"):
            clean = re.sub(r"```(?:json)?", "", clean).strip().rstrip("`").strip()
        return json.loads(clean)
    except Exception:
        return []

# ── FastAPI ───────────────────────────────────────────────────────
app = FastAPI(title="Cortex API")

@app.post("/api/upload")
async def upload(
    file: UploadFile = File(...),
    api_key: str = Form(...)
):
    if not api_key.strip():
        raise HTTPException(400, "API key required")
    raw = await file.read()
    filename = file.filename or "upload"
    ext = Path(filename).suffix.lower()
    client = claude_client(api_key)

    try:
        if ext == ".pdf":
            result = process_pdf(raw, filename, client)
            ftype = "pdf"
        elif ext in {".jpg", ".jpeg", ".png", ".webp", ".gif"}:
            result = process_image(raw, filename, client)
            ftype = "image"
        elif ext in {".mp4", ".mov", ".avi", ".mkv", ".webm"}:
            result = process_video(raw, filename, client)
            ftype = "video"
        else:
            result = process_text(raw, filename)
            ftype = "text"

        meta = generate_summary_and_entities(client, result["raw_text"], filename)
        chunks = result["chunks"]
        if chunks and HAS_EMBEDDER:
            vecs = embed_chunks(chunks)
            for i, c in enumerate(chunks):
                c["embedding"] = vecs[i].tolist()
        else:
            for c in chunks:
                c["embedding"] = [0.0] * 384

        KB["files"][filename] = {
            "type": ftype,
            "summary": meta.get("summary", ""),
            "entities": meta.get("key_entities", []),
            "topics": meta.get("topics", []),
            "key_sections": meta.get("key_sections", []),
            "important_facts": meta.get("important_facts", []),
            "file_type_detected": meta.get("file_type_detected", ""),
            "chunk_count": len(chunks),
            "thumb_b64": result.get("thumb_b64"),
            "preview": result.get("raw_text", "")[:5000],
        }
        old_chunks = [c for c in KB["chunks"] if c["source"] != filename]
        for c in chunks:
            c["embedding"] = np.array(c["embedding"], dtype="float32")
        KB["chunks"] = old_chunks + chunks
        rebuild_index()

        connections = find_connections(client)

        return JSONResponse({
            "filename": filename,
            "type": ftype,
            "summary": meta.get("summary", ""),
            "entities": meta.get("key_entities", []),
            "topics": meta.get("topics", []),
            "key_sections": meta.get("key_sections", []),
            "important_facts": meta.get("important_facts", []),
            "chunk_count": len(chunks),
            "total_chunks": len(KB["chunks"]),
            "connections": connections,
            "thumb_b64": result.get("thumb_b64"),
            "preview": result.get("raw_text", "")[:5000],
        })
    except Exception as e:
        err = str(e)
        if "401" in err or "invalid_api_key" in err or "auth" in err.lower():
            raise HTTPException(401, "Invalid API key")
        raise HTTPException(500, err)


@app.post("/api/ask")
async def ask(query: str = Form(...), api_key: str = Form(...)):
    if not api_key.strip():
        raise HTTPException(400, "API key required")
    if not KB["chunks"]:
        raise HTTPException(400, "No documents indexed yet")
    client = claude_client(api_key)

    # Detect broad "list all / summarize everything" queries
    broad_keywords = ["list all", "all subjects", "all topics", "all courses", "all codes",
                      "all units", "every subject", "complete list", "full list", "what are all",
                      "how many", "entire", "overview", "syllabus contains", "subjects in"]
    query_lower = query.lower()
    is_broad = any(kw in query_lower for kw in broad_keywords)

    if is_broad:
        # For broad queries: use file summaries + first chunk of each page
        # This gives a document-wide view without blowing token limits
        summary_context = []
        for fname, finfo in KB["files"].items():
            summary_context.append(
                f"[FILE: {fname}]\n"
                f"Summary: {finfo.get('summary','')}\n"
                f"Key Sections: {', '.join(finfo.get('key_sections', []))}\n"
                f"Topics: {', '.join(finfo.get('topics', []))}\n"
                f"Entities: {', '.join(finfo.get('key_entities', finfo.get('entities', [])))}"
            )
        # Also grab first chunk from every unique page for full coverage
        seen_pages = set()
        page_chunks = []
        for c in KB["chunks"]:
            key = (c["source"], c.get("page", 0))
            if key not in seen_pages:
                seen_pages.add(key)
                page_chunks.append(f"[{c['source']} p.{c.get('page',0)}] {c['text'][:300]}")
        context = "\n\n".join(summary_context) + "\n\nPER-PAGE EXCERPTS:\n" + "\n".join(page_chunks[:30])
        source_file = list(KB["files"].keys())[0] if KB["files"] else "document"
        source_hint = [{"file": source_file, "page": 0}]
    else:
        # For specific queries: use hybrid vector + keyword retrieval
        relevant = retrieve(query, top_k=10)
        context_parts = []
        for c in relevant[:10]:
            text = c["text"][:800]
            context_parts.append(f"[Source: {c['source']}, Page: {c.get('page',0)}]\n{text}")
        context = "\n\n".join(context_parts)
        source_hint = [{"file": c["source"], "page": c.get("page", 0)} for c in relevant[:3]]

    prompt = f"""You are Cortex, an intelligent document assistant. Answer the question using the provided context.
Be specific, complete and detailed. If the question asks to "list all" something, make sure you include EVERY item found across ALL pages.
Extract exact names, codes, titles from the context — do not summarize or omit items.
Only say "not found" if truly absent after checking all context.

Context:
{context[:6000]}

Question: {query}

Respond in JSON only (no markdown backticks):
{{"answer": "complete detailed answer here", "sources": [{{"file": "filename", "page": 0}}], "confidence": "high/medium/low"}}"""
    try:
        result = claude_text(client, prompt)
        clean = result.strip()
        if clean.startswith("```"):
            clean = re.sub(r"```(?:json)?", "", clean).strip().rstrip("`").strip()
        data = json.loads(clean)
        return JSONResponse(data)
    except Exception as e:
        return JSONResponse({"answer": f"Error: {str(e)}", "sources": [], "confidence": "low"})

@app.get("/api/status")
async def status():
    return JSONResponse({
        "files": {k: {**v, "thumb_b64": None} for k, v in KB["files"].items()},
        "total_chunks": len(KB["chunks"]),
        "file_count": len(KB["files"]),
    })

@app.post("/api/connections")
async def connections(api_key: str = Form(...)):
    if not api_key.strip():
        raise HTTPException(400, "API key required")
    client = claude_client(api_key)
    result = find_connections(client)
    return JSONResponse({"connections": result})

@app.delete("/api/file/{filename}")
async def delete_file(filename: str):
    if filename in KB["files"]:
        del KB["files"][filename]
        KB["chunks"] = [c for c in KB["chunks"] if c["source"] != filename]
        rebuild_index()
    return JSONResponse({"ok": True})

STATIC_DIR = Path(__file__).parent / "static"
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

@app.get("/")
async def root():
    return FileResponse(str(STATIC_DIR / "index.html"))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)