import asyncio
import hashlib
import logging
import re
import time
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
from typing import List, Tuple, Optional
import tempfile
import os
import platform

import faiss
import fitz
import google.generativeai as genai
import nltk
import numpy as np
import requests

from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# ------------------- Cross-Platform NLTK Setup -------------------

def setup_nltk():
    """Setup NLTK data directory based on platform and environment."""
    if os.getenv('NLTK_DATA'):
        nltk_data_dir = os.getenv('NLTK_DATA')
    else:
        system = platform.system().lower()
        if system == "windows":
            base_dir = os.getenv('LOCALAPPDATA', tempfile.gettempdir())
            nltk_data_dir = os.path.join(base_dir, "nltk_data")
        elif system == "darwin":  # macOS
            home = os.path.expanduser("~")
            base_dir = os.path.join(home, "Library", "Caches")
            if os.path.exists(base_dir):
                nltk_data_dir = os.path.join(base_dir, "nltk_data")
            else:
                nltk_data_dir = os.path.join(tempfile.gettempdir(), "nltk_data")
        else:  # Linux and others
            nltk_data_dir = os.path.join(tempfile.gettempdir(), "nltk_data")
    os.makedirs(nltk_data_dir, exist_ok=True)
    try:
        nltk.data.find("tokenizers/punkt")
    except LookupError:
        try:
            nltk.download("punkt", download_dir=nltk_data_dir, quiet=True)
        except Exception as e:
            logging.warning(f"Failed to download punkt to {nltk_data_dir}: {e}")
        try:
            nltk.download("punkt", quiet=True)
        except Exception:
            logging.error("Failed to download punkt data. Tokenization may fail.")
    if nltk_data_dir not in nltk.data.path:
        nltk.data.path.append(nltk_data_dir)
    return nltk_data_dir

# Setup NLTK
nltk_data_dir = setup_nltk()

try:
    from nltk.tokenize import sent_tokenize
except ImportError:
    logging.error("NLTK tokenizer not available. Install NLTK and download punkt data.")

def sent_tokenize(text):
    # fallback if nltk tokenizer unavailable
    import re
    sentences = re.split(r'[.!?]+', text)
    return [s.strip() for s in sentences if s.strip()]

# -------------- SPEED-OPTIMIZED Configuration --------------

GEMINI_API_KEY = "AIzaSyCYVp6yfDYGfUTGdbeuibEp5X5-WYfRDmg"
SECRET_TOKEN = "bc7aba495dfdaa88e718dec0d0fba29a2d45eaecac3268453778d027c7419081"

MAX_WORKERS = min(64, (os.cpu_count() or 4) * 4)  # Max threads for I/O
EMBED_BATCH = 400  # Batch size for embedding calls

CHUNK_WORDS = 600  # Chunk size in words
OVERLAP_WORDS = 30  # Small overlap to catch boundary context, test effect (was 0)

RETRIEVE_K = 4  # Increased from 3 to 4 for more context chunks

SIMILARITY_THRESHOLD = 0.2  # Increased threshold (from 0.15) to reduce noise

MAX_OUTPUT_TOKENS = 120  # Increased from 90 to 120 for richer answers

CONCURRENCY_LIMIT = min(12, MAX_WORKERS // 2)  # Tune per resource availability

genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel("gemini-1.5-flash")

executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("hackrx")

logger.info(f"Platform: {platform.system()} {platform.release()}")
logger.info(f"NLTK data directory: {nltk_data_dir}")
logger.info(f"Max workers: {MAX_WORKERS}, Concurrency limit: {CONCURRENCY_LIMIT}")

# ---------- SIMPLE FAST CACHE ----------

class FastCache:
    def __init__(self, max_size=1000):
        self.store = {}
        self.max_size = max_size

    def get(self, key: str):
        return self.store.get(key)

    def set(self, key: str, value):
        if len(self.store) >= self.max_size:
            # Remove oldest (simple FIFO)
            oldest_key = next(iter(self.store))
            self.store.pop(oldest_key)
        self.store[key] = value

    def clear(self):
        self.store.clear()

# Global caches
doc_cache = FastCache(50)
emb_cache = FastCache(3000)
ans_cache = FastCache(1000)

# ---------- FastAPI Setup ----------

app = FastAPI(title="HackRx LLM API Final", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    allow_credentials=True,
)

class Req(BaseModel):
    documents: str
    questions: List[str]

class Resp(BaseModel):
    answers: List[str]

async def verify_token(authorization: str = Header()):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(401, "Invalid authorization header")
    token = authorization.split()[1]
    if token != SECRET_TOKEN:
        raise HTTPException(401, "Invalid token")

# ---------- ULTRA-FAST PDF Extraction ----------

async def extract_text(url: str) -> str:
    key = hashlib.md5(url.encode()).hexdigest()
    cached = doc_cache.get(f"txt_{key}")
    if cached:
        return cached

    def _extract():
        try:
            r = requests.get(url, timeout=10)
            r.raise_for_status()
            with fitz.open(stream=r.content, filetype="pdf") as doc:
                text = " ".join(page.get_text() for page in doc if page.get_text().strip())
            return re.sub(r'\s+', ' ', text).strip()
        except Exception as e:
            logger.error(f"PDF extraction failed: {e}")
            raise HTTPException(500, f"Failed to extract PDF: {str(e)}")

    text = await asyncio.get_event_loop().run_in_executor(executor, _extract)

    if not text:
        raise HTTPException(400, "Extracted PDF text is empty")

    doc_cache.set(f"txt_{key}", text)
    return text

# ---------- FAST Chunking (with small overlap) ----------

def chunk_text(text: str) -> List[str]:
    words = text.split()
    chunks = []
    step = CHUNK_WORDS - OVERLAP_WORDS
    if step <= 0:
        step = CHUNK_WORDS  # fallback to no overlap if misconfigured
    for i in range(0, len(words), step):
        chunk = " ".join(words[i:i + CHUNK_WORDS])
        if chunk.strip():
            chunks.append(chunk)
    return chunks

# ---------- ULTRA-FAST Embeddings ----------

@lru_cache(maxsize=16384)  # Increased cache size from 8192
def embed_one(text: str, task: str):
    key = hashlib.md5((task + text).encode()).hexdigest()[:16]
    cached = emb_cache.get(key)
    if cached:
        return cached
    try:
        emb = genai.embed_content(model="models/embedding-001", content=text, task_type=task)["embedding"]
    except Exception as e:
        logger.error(f"Embedding error: {e}")
        emb = [0.0] * 768
    emb_cache.set(key, emb)
    return emb

async def embed_batch(texts: List[str], task: str) -> List[List[float]]:
    if not texts:
        return []
    sem = asyncio.Semaphore(20)  # High concurrency

    async def embed_single_async(text):
        async with sem:
            return await asyncio.get_event_loop().run_in_executor(executor, embed_one, text, task)

    return await asyncio.gather(*[embed_single_async(t) for t in texts])

# ---------- FAST FAISS Index ----------

def build_index(embeddings: List[List[float]]) -> Tuple[faiss.Index, np.ndarray]:
    if not embeddings:
        raise ValueError("No embeddings provided")
    mat = np.array(embeddings, dtype=np.float32)
    faiss.normalize_L2(mat)
    idx = faiss.IndexFlatIP(mat.shape[1])
    idx.add(mat)
    return idx, mat

# ---------- SOFT KEYWORD BOOSTING (OPTIONAL) ----------

def soft_keyword_boost(question: str, candidates: List[Tuple[float, str]]) -> List[Tuple[float, str]]:
    """
    Given a list of (similarity, chunk) tuples,
    boost similarity by a small factor if chunk contains question keywords.
    This is soft boosting, not strict filtering.
    """
    keywords = set(re.findall(r"\w+", question.lower()))
    boosted = []
    for sim, chunk in candidates:
        chunk_words = set(re.findall(r"\w+", chunk.lower()))
        common = keywords.intersection(chunk_words)
        boost_factor = 1.0 + 0.05 * len(common)  # 5% boost per keyword overlap, capped elsewhere
        boosted.append((sim * boost_factor, chunk))
    # Sort descending by boosted similarity score
    boosted.sort(key=lambda x: x[0], reverse=True)
    return boosted

# ---------- FAST Retrieval with soft keyword boosting ----------

async def retrieve(query: str, chunks: List[str], index: faiss.Index) -> List[str]:
    if not query or not chunks:
        return []
    q_emb = (await embed_batch([query], "retrieval_query"))[0]
    q_vec = np.array([q_emb], dtype=np.float32)
    faiss.normalize_L2(q_vec)
    top_k = min(RETRIEVE_K * 3, len(chunks))  # Search more initially for boosting

    D, I = index.search(q_vec, top_k)
    candidates = []
    for score, idx in zip(D[0], I[0]):
        if score > SIMILARITY_THRESHOLD:
            candidates.append((score, chunks[idx]))

    # Apply optional soft keyword boosting for ranking
    boosted = soft_keyword_boost(query, candidates)

    # Select top RETRIEVE_K chunks after boosting
    selected = [chunk for _, chunk in boosted[:RETRIEVE_K]]
    return selected

# ---------- FAST Answer Generation ----------

async def generate(question: str, chunks: List[str]) -> str:
    if not chunks:
        return "Information not found in the policy."

    cache_key = hashlib.md5((question + "".join(chunks)).encode()).hexdigest()[:16]
    cached = ans_cache.get(cache_key)
    if cached:
        return cached

    # Compose prompt with selected chunks
    context = "\n\n".join(chunks)
    prompt = f"""Context: {context}

Question: {question}

Answer briefly and precisely. If not found, say "Information not found in the policy."

Answer:"""

    def _gen():
        try:
            resp = model.generate_content(
                prompt,
                generation_config=genai.types.GenerationConfig(
                    temperature=0.1,
                    max_output_tokens=MAX_OUTPUT_TOKENS,
                    top_p=0.7,  # Lower for accuracy
                    top_k=20,
                ),
            )
            return re.sub(r'\s+', ' ', resp.text.strip()) or "Information not found in the policy."
        except Exception as e:
            logger.error(f"Generation error: {e}")
            return "Information not found in the policy."

    answer = await asyncio.get_event_loop().run_in_executor(executor, _gen)
    ans_cache.set(cache_key, answer)
    return answer

# ---------- MAIN PIPELINE ----------

async def process(doc_url: str, questions: List[str]) -> List[str]:
    start = time.time()
    text = await extract_text(doc_url)
    chunks = chunk_text(text)
    logger.info(f"Extracted and chunked in {time.time() - start:.2f}s ({len(chunks)} chunks)")

    idx_key = hashlib.md5("".join(chunks).encode()).hexdigest()[:16]
    cached_idx = doc_cache.get(f"idx_{idx_key}")
    if cached_idx:
        index, mat = cached_idx
        logger.info("Using cached index")
    else:
        emb_start = time.time()
        embeddings = await embed_batch(chunks, "retrieval_document")
        logger.info(f"Embeddings took {time.time() - emb_start:.2f}s")
        index, mat = build_index(embeddings)
        doc_cache.set(f"idx_{idx_key}", (index, mat))

    sem = asyncio.Semaphore(CONCURRENCY_LIMIT)

    async def handle_q(q):
        async with sem:
            ctx = await retrieve(q, chunks, index)
            return await generate(q, ctx)

    answers = await asyncio.gather(*[handle_q(q) for q in questions])
    total_time = time.time() - start
    logger.info(f"Total processing time: {total_time:.2f}s for {len(questions)} questions")
    return answers

# ---------- Endpoints ----------

@app.post("/hackrx/run", response_model=Resp)
async def run(req: Req, token=Depends(verify_token)):
    return {"answers": await process(req.documents, req.questions)}

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "platform": platform.system(),
        "python_version": platform.python_version(),
        "nltk_data_dir": nltk_data_dir,
        "cache_stats": {
            "doc_cache": len(doc_cache.store),
            "emb_cache": len(emb_cache.store),
            "ans_cache": len(ans_cache.store),
        }
    }

@app.get("/")
async def root():
    return {"message": "HackRx LLM API Final is running"}

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