import time
import uuid
from typing import List, Dict, Any
from concurrent.futures import ThreadPoolExecutor
import asyncio

from modules.ocr.receipt_detector import detect_receipt_type
from modules.ocr.models_donut import donut_model
from modules.ocr.models_paddle import paddle_model
from modules.ocr.models_trocr import trocr_model
from core.telemetry import log_telemetry

# Job storage (in production, use Redis or persistent DB)
jobs = {}

def get_model_priority(receipt_type: str) -> List[str]:
    """Get model priority based on receipt type"""
    priorities = {
        "structured": ["donut", "paddle", "trocr"],
        "printed": ["trocr", "donut", "paddle"],
        "handwritten": ["paddle", "trocr", "donut"],
        "blurry": ["paddle", "trocr", "donut"]
    }
    return priorities.get(receipt_type, ["donut", "paddle", "trocr"])

def process_single_receipt(image_path: str, filename: str) -> Dict[str, Any]:
    """Process one receipt through the cascade"""
    start_time = time.time()
    receipt_type = detect_receipt_type(image_path)
    
    models = get_model_priority(receipt_type)
    models_tried = []
    best_result = None
    best_confidence = 0.0
    best_model = ""
    
    for model_name in models:
        models_tried.append(model_name)
        
        if model_name == "donut":
            result, confidence = donut_model.load(), 0.0
            result, confidence = donut_model.process(image_path)
        elif model_name == "paddle":
            result, confidence = paddle_model.process(image_path)
        elif model_name == "trocr":
            result, confidence = trocr_model.process(image_path)
        else:
            continue
        
        if confidence >= best_confidence:
            best_confidence = confidence
            best_result = result
            best_model = model_name
        
        if confidence >= 0.70:
            break  # Stop cascade if good enough
    
    processing_time = int((time.time() - start_time) * 1000)
    
    # Log telemetry
    log_telemetry({
        "model_used": best_model,
        "models_tried": models_tried,
        "confidence": best_confidence,
        "receipt_type": receipt_type,
        "file_type": image_path.split('.')[-1],
        "field_count": len(best_result.keys()) if best_result else 0,
        "processing_ms": processing_time
    })
    
    warning = None
    if best_confidence < 0.70:
        warning = "No model reached 70% confidence. Best result shown."
    
    return {
        "filename": filename,
        "receiptType": receipt_type,
        "modelUsed": best_model,
        "modelsTried": models_tried,
        "overallConfidence": round(best_confidence, 2),
        "fields": [
            {"label": k.replace("_", " ").title(), "value": str(v), "confidence": best_confidence}
            for k, v in (best_result or {}).items() if k != "line_items"
        ],
        "lineItems": best_result.get("line_items", []) if best_result else [],
        "rawText": "",  # Would contain actual OCR text
        "processingMs": processing_time,
        "warning": warning
    }

async def process_batch_async(file_paths: List[tuple], job_id: str):
    """Process batch in background"""
    jobs[job_id] = {
        "status": "processing",
        "progress": 0,
        "total": len(file_paths),
        "results": []
    }
    
    loop = asyncio.get_event_loop()
    
    with ThreadPoolExecutor() as pool:
        for idx, (path, filename) in enumerate(file_paths):
            try:
                result = await loop.run_in_executor(
                    pool, process_single_receipt, path, filename
                )
                jobs[job_id]["results"].append(result)
                jobs[job_id]["progress"] = idx + 1
            except Exception as e:
                jobs[job_id]["results"].append({
                    "filename": filename,
                    "error": str(e),
                    "receiptType": "error",
                    "modelUsed": "none",
                    "modelsTried": [],
                    "overallConfidence": 0.0,
                    "fields": [],
                    "lineItems": [],
                    "rawText": "",
                    "processingMs": 0,
                    "warning": str(e)
                })
    
    jobs[job_id]["status"] = "completed"
    jobs[job_id]["progress"] = len(file_paths)

def create_job() -> str:
    job_id = str(uuid.uuid4())
    jobs[job_id] = {"status": "pending", "progress": 0, "total": 0, "results": []}
    return job_id

def get_job_status(job_id: str) -> dict:
    return jobs.get(job_id, {"status": "not_found"})

def get_job_results(job_id: str) -> List[Dict]:
    job = jobs.get(job_id)
    if job and job["status"] == "completed":
        return job["results"]
    return []