# Full LLM translation and processing combined with BM25 which helps idenitfy key words.
import pytesseract
from pdf2image import convert_from_path
import requests
import json
import time
import asyncio
import pypdf
import re
import hashlib
import base64
from io import BytesIO

ollama_semaphore = asyncio.Semaphore(1)
cloud_lock = asyncio.Lock()
LAST_CLOUD_CALL_TIME = 0.0
CLOUD_DELAY_SECONDS = 13.0  

def get_file_hash(filepath):
    h = hashlib.sha256()
    with open(filepath, 'rb') as file:
        while chunk := file.read(8192):
            h.update(chunk)
    return h.hexdigest()

def extract_text_from_pdf(doc_path):
    extracted_chunks = []
    try:
        reader = pypdf.PdfReader(doc_path)
        for page_idx, page in enumerate(reader.pages):
            text = page.extract_text()
            if text and len(text.strip()) > 50: 
                _chunk_text_layout_aware(text, page_idx, extracted_chunks)
            else:
                images = convert_from_path(doc_path, first_page=page_idx+1, last_page=page_idx+1)
                if images:
                    ocr_text = pytesseract.image_to_string(images[0], lang="eng+spa+ara+chi_sim+ita+msa")
                    _chunk_text_layout_aware(ocr_text, page_idx, extracted_chunks)
    except Exception as e:
        print(f"Extraction error: {e}")
    return extracted_chunks

def _chunk_text_layout_aware(text, page_idx, extracted_chunks):
    text = re.sub(r'\n{3,}', '\n\n', text)
    blocks = text.split('\n\n')
    current_chunk = ""
    for block in blocks:
        if len(current_chunk) + len(block) < 800:
            current_chunk += block + "\n\n"
        else:
            if current_chunk.strip():
                extracted_chunks.append({"page": page_idx + 1, "text": current_chunk.strip()})
            current_chunk = current_chunk[-150:] + block + "\n\n"
            
    if current_chunk.strip():
        extracted_chunks.append({"page": page_idx + 1, "text": current_chunk.strip()})

def get_page_image_b64(doc_path, page_num):
    images = convert_from_path(doc_path, first_page=page_num, last_page=page_num)
    if not images: return None
    buffered = BytesIO()
    images[0].save(buffered, format="JPEG", quality=85)
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

async def expand_query_async(query, document_anchor, api_key=""):
    """Dynamically detects document language via Anchor and expands the query."""
    sys_prompt = (
        "You are an expert multilingual legal translator. "
        "1. Read the document excerpt to detect its language and jurisdiction.\n"
        "2. Translate the 'Target Field' into 5 exact synonyms used in that specific language context.\n"
        "3. Output ONLY the synonyms separated by spaces. No other text."
    )
    user_prompt = f"Document Excerpt: {document_anchor[:1000]}\n\nTarget Field: {query}"
    
    try:
        if api_key:
            url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={api_key}"
            payload = {"contents": [{"parts": [{"text": f"{sys_prompt}\n\n{user_prompt}"}]}]}
            res = await asyncio.to_thread(requests.post, url, headers={"Content-Type": "application/json"}, json=payload, timeout=10)
            return res.json()['candidates'][0]['content']['parts'][0]['text'].strip()
        return query 
    except:
        return query

async def query_llm_async(model_name, system_prompt, user_prompt, response_type="extraction", api_key="", doc_path=None, page_nums=None):
    """CLOUD-ONLY ROUTER"""
    
    if not api_key: 
        return "Error: Gemini API Key Required", 0.0
        
    # Set schemas
    if response_type == "extraction":
        schema = {
            "type": "OBJECT", 
            "properties": {
                "step_1_evidence": {"type": "STRING"},
                "step_2_math_and_logic": {"type": "STRING"},
                "extracted_value": {"type": "STRING"}
            }, 
            "required": ["step_1_evidence", "step_2_math_and_logic", "extracted_value"]
        }
        target_key = "extracted_value"
    else:
        schema = {
            "type": "OBJECT", 
            "properties": {
                "internal_calculations_do_not_show_user": {"type": "STRING"},
                "final_response": {"type": "STRING"}
            }, 
            "required": ["internal_calculations_do_not_show_user", "final_response"]
        }
        target_key = "final_response"

    # Rate limiting lock
    global LAST_CLOUD_CALL_TIME
    async with cloud_lock:
        current_time = time.time()
        if current_time - LAST_CLOUD_CALL_TIME < CLOUD_DELAY_SECONDS:
            await asyncio.sleep(CLOUD_DELAY_SECONDS - (current_time - LAST_CLOUD_CALL_TIME))
        LAST_CLOUD_CALL_TIME = time.time()

    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={api_key}"
    parts = [{"text": f"SYSTEM INSTRUCTIONS:\n{system_prompt}\n\nUSER PROMPT:\n{user_prompt}"}]
    
    if doc_path and page_nums:
        parts[0]["text"] += "\n\nLook at the provided document images to find the exact value."
        for p_num in page_nums[:2]: 
            b64_img = await asyncio.to_thread(get_page_image_b64, doc_path, p_num)
            if b64_img: parts.append({"inlineData": {"mimeType": "image/jpeg", "data": b64_img}})
    
    payload = {
        "contents": [{"parts": parts}], 
        "generationConfig": {
            "temperature": 0.0,
            "responseMimeType": "application/json",
            "responseSchema": schema
        }
    }
    
    start_time = time.time()
    try:
        response = await asyncio.to_thread(requests.post, url, headers={"Content-Type": "application/json"}, json=payload, timeout=90)
        if response.status_code == 429: 
            return "RATE_LIMIT_EXCEEDED", round(time.time() - start_time, 2)
        response.raise_for_status()
        
        raw_content = response.json()['candidates'][0]['content']['parts'][0]['text']
        latency = round(time.time() - start_time, 2)
        try: 
            return str(json.loads(raw_content).get(target_key, "Not Found")).strip(), latency
        except: 
            return "Parse Error", latency
    except Exception as e: 
        return f"API Error", round(time.time() - start_time, 2)