from typing import List, Dict, Any
from rag.ocr import extract_text as ocr_extract_text
from rag.chunker import chunk_text as chunker_chunk_text
from rag.retriever import get_pinecone_store as retriever_get_pinecone_store, retrieve_hybrid
from rag.reranker import rerank
from rag.engine import agent_stream as engine_agent_stream

# Re-expose the system prompt for any direct imports (though now managed in rag/prompt.py)
from rag.prompt import SYSTEM_PROMPT

def extract_text(file_bytes: bytes) -> List[Dict[str, Any]]:
    """
    Extracts text from PDF pages (Proxy to rag/ocr.py).
    """
    return ocr_extract_text(file_bytes)

def chunk_text(pages: List[Dict[str, Any]], filename: str) -> List[Dict[str, Any]]:
    """
    Splits pages into recursive/semantic chunks (Proxy to rag/chunker.py).
    """
    return chunker_chunk_text(pages, filename)

def save_to_pinecone(machine_id: str, chunks: List[Dict[str, Any]]):
    """
    Processes raw chunks (classification, metadata generation, sequential linking)
    and saves them to both Pinecone and the local cache (Proxy to rag/ingestion.py logic).
    """
    if not chunks:
        return
        
    filename = chunks[0]["filename"]
    # Use first chunk text as sample for classification
    sample_text = chunks[0]["text"] if chunks else ""
    
    from rag.classifier import classify_document
    from rag.metadata import generate_metadata
    from langchain_core.documents import Document
    from config import STORE_DIR
    import json
    
    # 1. Automatically classify document type
    doc_type = classify_document(filename, sample_text)
    print(f"RAG Proxy: Classified '{filename}' as '{doc_type}'")
    
    # 2. Generate heading detection and sequential links
    enriched_chunks = generate_metadata(chunks, machine_id, doc_type)
    
    # 3. Upload to Pinecone
    try:
        vector_store = retriever_get_pinecone_store(machine_id)
        documents = [
            Document(
                page_content=c["text"],
                metadata={
                    "filename": c["filename"],
                    "page": c["page"],
                    "chunk_id": c["chunk_id"],
                    "document_type": c["document_type"],
                    "section": c["section"],
                    "heading": c["heading"],
                    "section_number": c.get("section_number", ""),
                    "previous_chunk": c["previous_chunk"] or "",
                    "next_chunk": c["next_chunk"] or "",
                    "machine_id": machine_id
                }
            )
            for c in enriched_chunks
        ]
        vector_store.add_documents(documents)
        print(f"RAG Proxy: Successfully uploaded {len(documents)} chunks to Pinecone.")
    except Exception as e:
        print(f"RAG Proxy: Failed to upload to Pinecone: {e}")
        
    # 4. Save to local chunks cache (for BM25 indexing & context expansion)
    from rag.retriever import load_local_chunks, clear_bm25_cache
    local_chunks = load_local_chunks(machine_id)
    # Remove existing chunks for this file if we are re-uploading
    local_chunks = [c for c in local_chunks if c["filename"] != filename]
    local_chunks.extend(enriched_chunks)
    
    cache_dir = STORE_DIR / machine_id
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = cache_dir / "chunks_cache.json"
    with open(cache_path, "w", encoding="utf-8") as f:
        json.dump(local_chunks, f, indent=2, ensure_ascii=False)

    try:
        from database import save_chunks_to_db
        save_chunks_to_db(machine_id, enriched_chunks)
    except Exception as e:
        print(f"RAG Proxy: Failed to save chunks to PostgreSQL: {e}", flush=True)

    # 5. Clear the BM25 cache for this machine to force rebuild on next query
    clear_bm25_cache(machine_id)

    # 6. Invalidate cached chat answers for this machine — a cached "not found" or a
    # cached answer built from the old chunk set must not outlive this new upload.
    from rag.engine import clear_query_cache
    clear_query_cache(machine_id)


def get_pinecone_store(machine_id: str):
    """
    Retrieves the Pinecone Vector Store for the machine namespace.
    """
    return retriever_get_pinecone_store(machine_id)

def retrieve(query: str, machine_id: str) -> List[Dict[str, Any]]:
    """
    Retrieves, reranks, and expands adjacent context in correct reading order.
    Used by the AI task and alert extraction endpoints in main.py.
    """
    from rag.classifier import classify_question
    from rag.query_expander import expand_query
    from rag.retriever import expand_context
    
    # 1. Classify query intent
    intent = classify_question(query)
    
    # 2. Expand queries
    queries = expand_query(query)
    
    # 3. Perform hybrid retrieval
    merged = retrieve_hybrid(queries, machine_id, intent)
    if not merged:
        return []
        
    # 4. Rerank to top 12 (increased recall to capture entire tables)
    top_reranked = rerank(query, merged, top_n=12)
    
    # 5. Run context expansion to pull adjacent table parts and sort them in correct page order
    expanded = expand_context(top_reranked, machine_id)
    
    # 6. Map back to the exact format main.py expects (in correct reading sequence)
    results = []
    for idx, c in enumerate(expanded):
        results.append({
            "id": idx,
            "text": c["text"],
            "filename": c["filename"],
            "score": c.get("score", 0.0),
            "page": c["page"]
        })
    return results

async def agent_stream(message: str, machine_id: str, history: List[Dict[str, Any]]):
    """
    Streams RAG v2 responses using SSE (Proxy to rag/engine.py).
    """
    async for event in engine_agent_stream(message, machine_id, history):
        yield event