import os
import sys
import hashlib
import time
import asyncio
import httpx

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ingestion.ocr_processor import download_and_extract

# --- SSOT INGESTION CONFIGURATION ---
HF_SPACE_URL = os.environ.get("HF_SPACE_URL", "https://harshrawat18-govbridge-api.hf.space")
ADMIN_SECRET = os.environ.get("ADMIN_SECRET", "")
PDF_URL = os.environ.get("PDF_URL", "")

def chunk_text(text: str, chunk_size: int = 300) -> list:
    """Split text into word-based chunks with overlap."""
    words = text.split()
    chunks = []
    overlap = 30
    step = chunk_size - overlap
    for i in range(0, len(words), step):
        chunk = " ".join(words[i:i + chunk_size])
        if chunk.strip():
            chunks.append(chunk)
    return chunks

async def ingest_chunk_async(
    client: httpx.AsyncClient,
    chunk: str,
    metadata: dict,
    chunk_index: int,
    semaphore: asyncio.Semaphore
) -> bool:
    """Ingest a single chunk with semaphore-controlled concurrency."""
    async with semaphore:
        content_hash = hashlib.sha256(
            f"{metadata['source_url']}_{chunk_index}".encode()
        ).hexdigest()
        
        payload = {
            "title": metadata.get("title", "Government Document"),
            "text": chunk,
            "source_url": metadata["source_url"],
            "ministry": metadata.get("ministry", "Government of India"),
            "state": metadata.get("state", "National"),
            "doc_type": metadata.get("doc_type", "document"),
            "content_hash": content_hash,
            "chunk_index": chunk_index
        }
        
        try:
            response = await client.post(
                f"{HF_SPACE_URL}/api/admin/ingest",
                json=payload,
                headers={"X-Admin-Secret": ADMIN_SECRET},
                params={"admin_key": ADMIN_SECRET},
                timeout=30.0
            )
            if response.status_code == 200:
                print(f"✅ Chunk {chunk_index+1} injected instantly")
                return True
            else:
                print(f"⚠️ Chunk {chunk_index+1} failed: {response.status_code}")
                return False
        except Exception as e:
            print(f"❌ Chunk {chunk_index+1} error: {e}")
            return False

async def ingest_all_chunks_async(chunks: list, metadata: dict):
    """Dispatch all chunks concurrently with controlled batch size."""
    semaphore = asyncio.Semaphore(10)
    async with httpx.AsyncClient() as client:
        tasks = [
            ingest_chunk_async(client, chunk, metadata, i, semaphore)
            for i, chunk in enumerate(chunks)
        ]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
    successful = sum(1 for r in results if r is True)
    failed = len(results) - successful
    print(f"✅ Ingested {successful}/{len(chunks)} chunks")
    if failed > 0:
        print(f"⚠️ {failed} chunks failed — retry on next run")
    return successful

def main():
    if not PDF_URL:
        print("❌ No PDF_URL provided in terminal")
        sys.exit(1)

    if not ADMIN_SECRET:
        print("❌ 🔒 SECURITY HALT: ADMIN_SECRET environment variable is missing!")
        sys.exit(1)

    print(f"🚀 Starting Async PDF extraction: {PDF_URL}")

    try:
        result = download_and_extract(PDF_URL)
        text = result["text"]
        word_count = len(text.split())
        print(f"✅ Extracted {word_count} words in memory via {result['extraction_method']}")
        
        chunks = chunk_text(text, chunk_size=300)
        print(f"📦 Sliced into {len(chunks)} high-density AI chunks")

        # Start precise benchmarking timer
        start_time = time.time()
        
        asyncio.run(ingest_all_chunks_async(chunks, {
            "source_url": PDF_URL,
            "title": "Government Policy Document",
            "ministry": "Government of India",
            "state": "National",
            "doc_type": "document"
        }))

        # End timer
        elapsed_time = time.time() - start_time
        print(f"⏱️ Async Ingestion Time: {elapsed_time:.2f} seconds")
        print(f"🎉 MASSIVE W! {len(chunks)} chunks ingested to Vector DB.")

    except Exception as e:
        print(f"❌ Critical failure: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()