import os
import uuid
import json
import concurrent.futures
from typing import List, Dict
from unstructured.partition.pdf import partition_pdf
from unstructured.chunking.title import chunk_by_title
from langchain_core.documents import Document
from langchain_pinecone import PineconeVectorStore
from langchain_core.messages import HumanMessage
from pinecone import Pinecone, ServerlessSpec

from config import *

# Initialize Pinecone Client
pc = Pinecone(api_key=PINECONE_API_KEY)

class SessionDocStore:
    """
    In-memory storage for heavy content (Images/Tables) tied to a session.
    In a real production app, use Redis/Postgres. 
    """
    def __init__(self):
        self.store = {}

    def save_chunk(self, doc_id: str, data: Dict):
        self.store[doc_id] = data

    def get_chunk(self, doc_id: str):
        return self.store.get(doc_id, {})
    
    def clear(self):
        self.store = {}
def cleanup_session_index(session_id: str):
    """
    Deletes the namespace from Pinecone. 
    Returns True if successful, raises Exception if failed.
    """
    if not session_id:
        print("⚠️ No session ID to clean.")
        return False
        
    print(f"🧹 Attempting to delete namespace: {session_id}")
    
    try:
        # Re-initialize index to ensure connection is fresh
        index = pc.Index(INDEX_NAME)
        
        # Check if index actually exists (Prevent silent failures)
        stats = index.describe_index_stats()
        
        # Execute Delete
        index.delete(delete_all=True, namespace=session_id)
        
        print(f"✅ Successfully deleted namespace: {session_id}")
        return True
        
    except Exception as e:
        error_msg = f"❌ Pinecone Delete Failed: {str(e)}"
        print(error_msg)
        raise Exception(error_msg) # Raise so App can see it

# Helper to summarize visual content for embeddings
def create_multimodal_summary(text, tables, images):
    llm = get_llm()
    prompt_text = f"Analyze content. TEXT: {text[:1000]}. INSTRUCTIONS: Summarize text and describe images/tables for retrieval."
    
    message_content = [{"type": "text", "text": prompt_text}]
    if images:
        for b64_str in images:
            if "," in b64_str: b64_str = b64_str.split(",")[1]
            message_content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64_str}", "detail": "low"}
            })
            
    response = llm.invoke([HumanMessage(content=message_content)])
    return response.content

def process_single_chunk(i, chunk, doc_store):
    """
    Worker function to process a single chunk in a separate thread.
    """
    content = {'text': chunk.text, 'tables': [], 'images': []}
    
    # Extract visual data
    if hasattr(chunk, 'metadata') and hasattr(chunk.metadata, 'orig_elements'):
        for element in chunk.metadata.orig_elements:
            el_type = type(element).__name__
            if el_type == 'Table':
                content['tables'].append(getattr(element.metadata, 'text_as_html', element.text))
            elif el_type == 'Image' and hasattr(element.metadata, 'image_base64'):
                # OPTIONAL: Add Image Size filtering here to skip small icons
                content['images'].append(element.metadata.image_base64)

    # Deciding whether to call LLM (Slow) or just use text (Fast)
    if content['images'] or content['tables']:
        # This is the bottleneck we are parallelizing
        enhanced_text = create_multimodal_summary(content['text'], content['tables'], content['images'])
    else:
        enhanced_text = content['text']

    doc_id = str(uuid.uuid4())
    
    # Save heavy data to local store
    doc_store.save_chunk(doc_id, {
        "raw_text": content['text'],
        "tables": content['tables'],
        "images": content['images']
    })
    
    # Return the processed Document
    return Document(
        page_content=enhanced_text,
        metadata={"doc_id": doc_id, "chunk_index": i}
    )

def process_and_ingest(file_path: str, session_id: str, doc_store: SessionDocStore):
    print(f"📄 Partitioning: {file_path}")
    
    # ⚡ OPTIMIZATION #1: Changed strategy from "hi_res" to "auto"
    # "auto" is much faster. Use "hi_res" only if "auto" fails to read tables.
    elements = partition_pdf(
        filename=file_path, 
        strategy="auto", 
        infer_table_structure=True,
        extract_image_block_types=["Image", "Table"], 
        extract_image_block_to_payload=True
    )
    
    chunks = chunk_by_title(elements, max_characters=2000, new_after_n_chars=1500, combine_text_under_n_chars=300)
    
    documents = []
    
    print(f"🔄 Processing {len(chunks)} chunks in parallel...")
    
    # ⚡ OPTIMIZATION #2: ThreadPoolExecutor for Parallel Summarization
    # We use 5 workers. Going higher might hit Azure Rate Limits (429 Errors).
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        # Submit all tasks
        future_to_chunk = {
            executor.submit(process_single_chunk, i, chunk, doc_store): i 
            for i, chunk in enumerate(chunks)
        }
        
        # Gather results as they complete (or strictly in order)
        # We iterate over the original range to maintain order
        futures_list = list(future_to_chunk.keys())
        
        for future in futures_list:
            try:
                doc = future.result() # Blocks until this specific chunk is done
                documents.append(doc)
            except Exception as e:
                print(f"❌ Error processing chunk: {e}")

    print(f"🔮 Ingesting {len(documents)} vectors to Namespace: {session_id}")
    
    PineconeVectorStore.from_documents(
        documents=documents,
        index_name=INDEX_NAME,
        embedding=get_embeddings(),
        namespace=session_id 
    )
    
    return documents # Return docs for BM25 initialization

def cleanup_session_index(session_id: str):
    """Deletes the specific namespace for the session"""
    try:
        index = pc.Index(INDEX_NAME)
        index.delete(delete_all=True, namespace=session_id)
        print(f"🗑️ Deleted Namespace: {session_id}")
    except Exception as e:
        print(f"⚠️ Error deleting namespace: {e}")