import os
# Fix Huggingface cache permissions for Huggingface Spaces and Docker
os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache"
os.environ["HF_HOME"] = "/tmp/.cache"
os.environ["HF_CACHE_DIR"] = "/tmp/.cache"

# Rest of your imports and code...
import asyncio
import hashlib
import logging
import re
import time
import tempfile
import platform
from typing import List
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
import faiss
import fitz
import nltk
import numpy as np
import requests
import google.generativeai as genai
from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer

# ... (YOUR WHOLE CODE AS BEFORE)


# For language detection; add to requirements.txt: langdetect
try:
    from langdetect import detect, DetectorFactory
    DetectorFactory.seed = 0
    LANGDETECT_AVAILABLE = True
except ImportError:
    LANGDETECT_AVAILABLE = False

def setup_nltk():
    if os.getenv("NLTK_DATA"):
        d = os.getenv("NLTK_DATA")
    else:
        sys_name = platform.system().lower()
        if sys_name == "windows":
            base = os.getenv("LOCALAPPDATA", tempfile.gettempdir())
        elif sys_name == "darwin":
            home = os.path.expanduser("~")
            base = os.path.join(home, "Library", "Caches")
            if not os.path.exists(base):
                base = tempfile.gettempdir()
        else:
            base = tempfile.gettempdir()
        d = os.path.join(base, "nltk_data")
    os.makedirs(d, exist_ok=True)
    for pkg in ["punkt_tab", "punkt"]:
        try:
            nltk.data.find(f"tokenizers/{pkg}")
            break
        except LookupError:
            try:
                nltk.download(pkg, download_dir=d, quiet=True)
            except:
                pass
    if d not in nltk.data.path:
        nltk.data.path.append(d)
    return d

nltk_data_dir = setup_nltk()

try:
    from nltk.tokenize import sent_tokenize
except ImportError:
    def sent_tokenize(text: str):
        return [s.strip() for s in re.split(r"[.!?]+", text) if s.strip()]

def detect_question_language(text: str) -> str:
    if not LANGDETECT_AVAILABLE:
        if any('\u0d00' <= char <= '\u0d7f' for char in text):
            return "ml"
        else:
            return "en"
    try:
        lang_code = detect(text)
        return lang_code
    except:
        return "en"

def get_language_name(lang_code: str) -> str:
    lang_map = {
        "en": "English", "ml": "Malayalam", "hi": "Hindi", "ar": "Arabic",
        "fr": "French", "es": "Spanish", "de": "German",
        "zh": "Chinese", "ja": "Japanese", "ko": "Korean"
    }
    return lang_map.get(lang_code, "English")

def get_not_found_message(lang_code: str) -> str:
    messages = {
        "en": "Information not found in the policy.",
        "ml": "പോളിസിയിൽ വിവരങ്ങൾ കണ്ടെത്തിയില്ല.",
        "hi": "पॉलिसी में जानकारी नहीं मिली।",
        "ar": "لم يتم العثور على المعلومات في البوليصة.",
        "fr": "Informations introuvables dans la police.",
        "es": "Información no encontrada en la póliza.",
        "de": "Informationen in der Police nicht gefunden.",
        "zh": "在保单中未找到信息。",
        "ja": "ポリシーに情報が見つかりません。",
        "ko": "정책에서 정보를 찾을 수 없습니다."
    }
    return messages.get(lang_code, "Information not found in the policy.")

def is_compound_question(question: str) -> bool:
    """Check if question has multiple parts that need separate answers"""
    compound_indicators = [
        " and what ", " and how ", " and why ", " and when ", " and where ",
        " as well as ", " also what ", " in addition ", "? what ", "? how ",
        "; what ", "; how ", "; why "
    ]
    return any(indicator in question.lower() for indicator in compound_indicators)

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "AIzaSyCYVp6yfDYGfUTGdbeuibEp5X5-WYfRDmg")
SECRET_TOKEN   = os.getenv("SECRET_TOKEN",   "bc7aba495dfdaa88e718dec0d0fba29a2d45eaecac3268453778d027c7419081")

MAX_WORKERS      = min(64, (os.cpu_count() or 4) * 4)
CHUNK_WORDS      = 150  # Slightly larger for more context
OVERLAP_WORDS    = 50   # More overlap for better coverage
PASS1_K          = 35   # Even more candidates
PASS2_K          = 15   # More final chunks
SIMILARITY_THRESHOLD = 0.08
MAX_OUTPUT_TOKENS    = 400  # Significantly increased for compound questions
CONCURRENCY_LIMIT    = 16
QUERY_VARIANTS       = 5    # More query variants

genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel("gemini-1.5-flash")
embedding_model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("hackrx")
logger.info(f"Platform: {platform.system()}  NLTK data: {nltk_data_dir}")

app = FastAPI(title="HackRx LLM API Multilingual Enhanced", version="7.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")
    if authorization.split()[1] != SECRET_TOKEN:
        raise HTTPException(401, "Invalid token")

async def extract_text(url: str) -> str:
    def _extract():
        r = requests.get(url, timeout=30)
        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()
    text = await asyncio.get_event_loop().run_in_executor(executor, _extract)
    if not text:
        raise HTTPException(400, "No text extracted from PDF.")
    return text

def chunk_text(text: str) -> List[str]:
    try:
        sentences = sent_tokenize(text)
        chunks = []
        current_chunk = []
        current_len = 0
        for sentence in sentences:
            sentence_words = sentence.split()
            if current_len + len(sentence_words) > CHUNK_WORDS and current_chunk:
                chunks.append(" ".join(current_chunk))
                current_chunk = current_chunk[-OVERLAP_WORDS:] if current_len >= OVERLAP_WORDS else []
                current_len = len(current_chunk)
            current_chunk.extend(sentence_words)
            current_len += len(sentence_words)
        if current_chunk:
            chunks.append(" ".join(current_chunk))
    except Exception as e:
        logger.warning(f"Sentence tokenization failed: {e}, fallback chunking")
        words = text.split()
        chunks = []
        step = CHUNK_WORDS - OVERLAP_WORDS
        for i in range(0, len(words), step):
            chunk = " ".join(words[i:i + CHUNK_WORDS])
            if chunk.strip():
                chunks.append(chunk)
    return [chunk for chunk in chunks if chunk.strip()]

def embed_one(text: str, task: str):
    return embedding_model.encode(text, convert_to_numpy=True, normalize_embeddings=True).tolist()

async def embed_batch(texts: List[str], task: str = "retrieval") -> List[List[float]]:
    sem = asyncio.Semaphore(20)
    async def process_one(t):
        async with sem:
            return await asyncio.get_event_loop().run_in_executor(executor, embed_one, t, task)
    return await asyncio.gather(*[process_one(t) for t in texts])

def build_index(embeddings: List[List[float]]) -> faiss.Index:
    mat = np.array(embeddings, dtype=np.float32)
    faiss.normalize_L2(mat)
    index = faiss.IndexFlatIP(mat.shape[1])
    index.add(mat)
    return index

def soft_keyword_boost(query: str, candidates: List[tuple]) -> List[tuple]:
    keywords = set(re.findall(r"\w+", query.lower()))
    boosted = []
    for sim, chunk in candidates:
        overlap = len(keywords.intersection(set(re.findall(r"\w+", chunk.lower()))))
        boosted.append((sim * (1.0 + 0.08 * overlap), chunk))
    boosted.sort(key=lambda x: x[0], reverse=True)
    return boosted

async def expand_question(question: str) -> List[str]:
    lang_code = detect_question_language(question)
    lang_name = get_language_name(lang_code)
    
    # For compound questions, create variants that focus on different parts
    if is_compound_question(question):
        prompt = f"""Generate {QUERY_VARIANTS} different phrasings and sub-questions for the following compound question in {lang_name}. 
Include variants that focus on different parts of the question:

Question: {question}

Generate alternative phrasings and focused sub-questions in {lang_name}:
1."""
    else:
        prompt = f"""Generate {QUERY_VARIANTS} different phrasings of the following question in {lang_name}, preserving meaning but varying wording:

Question: {question}

Generate alternative phrasings in {lang_name}:
1."""

    def _expand():
        try:
            resp = model.generate_content(
                prompt,
                generation_config=genai.types.GenerationConfig(
                    temperature=0.7, top_p=0.9, top_k=40, max_output_tokens=150
                ),
            )
            lines = [l.strip("1234567890.-• ") for l in resp.text.split("\n") if l.strip()]
            variants = list(set([question] + [l for l in lines if len(l) > 5]))[:QUERY_VARIANTS]
            return variants
        except Exception as e:
            logger.warning(f"Query expansion failed: {e}")
            return [question]
    return await asyncio.get_event_loop().run_in_executor(executor, _expand)

async def dual_pass_retrieve(query: str, chunks: List[str], index: faiss.Index) -> List[str]:
    variants = await expand_question(query)
    all_scores = []
    for var in variants:
        q_emb = (await embed_batch([var], "retrieval_query"))[0]
        q_vec = np.array([q_emb], dtype=np.float32)
        faiss.normalize_L2(q_vec)
        D, I = index.search(q_vec, min(PASS1_K, len(chunks)))
        for dist, idx in zip(D[0], I[0]):
            if dist > SIMILARITY_THRESHOLD:
                all_scores.append((float(dist), chunks[idx]))
    
    logger.info(f"Retrieved {len(all_scores)} candidate chunks for question: {query[:100]}")
    
    boosted = soft_keyword_boost(query, all_scores)
    seen = set()
    final_chunks = []
    for _, chunk in boosted:
        if chunk not in seen:
            seen.add(chunk)
            final_chunks.append(chunk)
            if len(final_chunks) >= PASS2_K:
                break
    
    logger.info(f"Final selected {len(final_chunks)} chunks for: {query[:70]}")
    return final_chunks

async def generate(question: str, chunks: List[str]) -> str:
    lang_code = detect_question_language(question)
    lang_name = get_language_name(lang_code)
    
    if not chunks:
        return get_not_found_message(lang_code)

    context = "\n\n".join(chunks)
    
    # Enhanced prompt for compound questions
    if is_compound_question(question):
        prompt = f"""You are an expert document analyst. Answer ALL PARTS of the compound question completely and thoroughly based ONLY on the provided context.

CRITICAL INSTRUCTIONS:
- This question has multiple parts - answer EVERY part completely
- Answer in {lang_name} language ONLY (the same language as the question)
- Be comprehensive and detailed for each part of the question
- Structure your answer to address each component clearly
- If any part cannot be answered from context, state that specifically
- Question language: {lang_name}
- Your response MUST be in {lang_name}

Context:
{context}

Compound Question: {question}

Provide a complete answer addressing ALL parts of the question in {lang_name}:"""
    else:
        prompt = f"""You are an expert document analyst. Answer the question comprehensively based ONLY on the provided context.

CRITICAL INSTRUCTIONS:
- Answer in {lang_name} language ONLY (the same language as the question)
- Be precise, detailed, and factual
- If information is not found in context, say: "{get_not_found_message(lang_code)}"
- Question language: {lang_name}
- Your response MUST be in {lang_name}

Context:
{context}

Question: {question}

Answer comprehensively in {lang_name}:"""

    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.8, 
                    top_k=25
                ),
            )
            answer = re.sub(r'\s+', ' ', resp.text.strip())
            return answer if answer else get_not_found_message(lang_code)
        except Exception as e:
            logger.error(f"Generation error: {e}")
            return get_not_found_message(lang_code)
    
    return await asyncio.get_event_loop().run_in_executor(executor, _gen)

async def process(doc_url: str, questions: List[str]) -> List[str]:
    # <==== LOGGING ADDED, safe and fast ====
    logger.info(f"Received request. PDF URL: {doc_url}")
    logger.info(f"Questions received: {questions}")
    # <==== END LOGGING ADDITION ====
    start_time = time.time()
    text = await extract_text(doc_url)
    chunks = chunk_text(text)
    logger.info(f"Created {len(chunks)} chunks")
    
    embeddings = await embed_batch(chunks, "retrieval_document")
    index = build_index(embeddings)
    logger.info(f"Built index with {len(embeddings)} embeddings")
    
    sem = asyncio.Semaphore(CONCURRENCY_LIMIT)
    
    async def handle_q(q):
        async with sem:
            lang_code = detect_question_language(q)
            lang_name = get_language_name(lang_code)
            is_compound = is_compound_question(q)
            
            logger.info(f"Processing {'compound' if is_compound else 'simple'} question in {lang_name}: {q[:80]}...")
            ctx = await dual_pass_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_time
    logger.info(f"Processed {len(questions)} questions in {total_time:.2f}s")
    
    return answers

@app.post("/hackrx/run", response_model=Resp)
async def run(req: Req, token=Depends(verify_token)):
    try:
        return {"answers": await process(req.documents, req.questions)}
    except Exception as e:
        logger.exception("Error in /hackrx/run:")
        raise HTTPException(500, f"Internal error: {str(e)}")

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "platform": platform.system(),
        "multilingual": True,
        "compound_questions": True,
        "langdetect_available": LANGDETECT_AVAILABLE,
        "embedding_model": "paraphrase-multilingual-MiniLM-L12-v2"
    }

@app.get("/")
async def root():
    return {
        "message": "HackRx Multilingual API Enhanced is running",
        "features": [
            "Automatic language detection",
            "Answer in question's language", 
            "Compound question handling",
            "Enhanced context retrieval",
            "Comprehensive answers"
        ]
    }

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