"""
Sequential Local Benchmark Runner.
Executes all 6 OCR models sequentially, captures accurate inference times,
generates bounding box visualizations, and computes diff & consensus metrics.
"""

import os
import time
import logging
from typing import Dict, Any, List, Optional
from PIL import Image

from core.models import OCRModelOutput, Region, ALL_REGION_TYPES
from core.region_classifier import compute_region_summary
from core.visualizer import save_visualization, image_to_base64_jpeg
from adapters import AVAILABLE_MODELS, MODEL_CATALOG, get_adapter_by_name
from utils.memory_manager import model_manager
from utils.diff_engine import generate_consensus_matrix, compute_diff, compute_similarity_ratio

logger = logging.getLogger("BenchmarkRunner")


def run_sequential_local_benchmark(
    image: Image.Image,
    output_dir: str = "outputs/visualizations",
    document_name: str = "document.png"
) -> Dict[str, Any]:
    """
    Executes all 6 models sequentially on local hardware,
    generating bounding box annotated images, performance metrics, and diff matrices.
    """
    os.makedirs(output_dir, exist_ok=True)
    benchmark_start = time.perf_counter()
    
    results: List[Dict[str, Any]] = []
    valid_texts: Dict[str, str] = {}

    img_w, img_h = image.size
    logger.info(f"Starting sequential local benchmark on '{document_name}' ({img_w}x{img_h} px)")

    for name in AVAILABLE_MODELS:
        logger.info(f"==> Benchmarking Model: {name}")
        adapter = get_adapter_by_name(name)
        
        # Execute adapter inference with timing & error capture
        res = adapter.process(image)
        
        is_success = (res.get("status") == "SUCCESS")
        regions_data: List[Region] = res.get("regions", [])
        
        # Convert dictionary regions to Region objects if needed
        parsed_regions: List[Region] = []
        for r in regions_data:
            if isinstance(r, Region):
                parsed_regions.append(r)
            elif isinstance(r, dict):
                parsed_regions.append(Region(
                    box=r.get("box", [0, 0, 0, 0]),
                    text=r.get("text", ""),
                    region_type=r.get("region_type", "Text"),
                    confidence=r.get("confidence"),
                    details=r.get("details")
                ))

        # Compute word count & region distribution
        raw_text = res.get("text") or res.get("markdown") or ""
        word_count = len(raw_text.split()) if raw_text else 0
        region_counts = compute_region_summary(parsed_regions) if is_success else {}

        # Render Annotated Visualization Image
        annotated_path = os.path.join(output_dir, f"{name.lower().replace('-', '_')}_annotated.png")
        annotated_base64 = None

        if is_success and parsed_regions:
            try:
                save_visualization(image, parsed_regions, annotated_path)
                annotated_img = Image.open(annotated_path)
                annotated_base64 = image_to_base64_jpeg(annotated_img, quality=85)
                logger.info(f"Rendered {len(parsed_regions)} bounding boxes for {name} -> {annotated_path}")
            except Exception as ve:
                logger.warning(f"Visualization rendering error for {name}: {ve}")

        output_item = OCRModelOutput(
            model_name=name,
            model_id=MODEL_CATALOG[name]["hf_model_id"],
            status=res.get("status", "ERROR"),
            inference_time_seconds=res.get("inference_time_seconds"),
            inference_time_str=res.get("inference_time_str", "N/A"),
            text=res.get("text"),
            markdown=res.get("markdown"),
            json=res.get("json"),
            output_type=res.get("output_type", MODEL_CATALOG[name]["preferred_output"]),
            word_count=word_count,
            regions=parsed_regions,
            region_counts=region_counts,
            annotated_image_path=annotated_path if (is_success and parsed_regions) else None,
            annotated_image_base64=annotated_base64,
            error=res.get("error")
        )

        results.append(output_item.to_dict())

        if is_success and raw_text.strip():
            valid_texts[name] = raw_text.strip()

    total_time = round(time.perf_counter() - benchmark_start, 2)
    logger.info(f"Completed local benchmark across 6 models in {total_time}s")

    # Pairwise Consensus Matrix on successful outputs only
    consensus_data = generate_consensus_matrix(valid_texts)

    return {
        "success": True,
        "document_name": document_name,
        "image_dimensions": {"width": img_w, "height": img_h},
        "total_benchmark_time_seconds": total_time,
        "available_region_types": ALL_REGION_TYPES,
        "results": results,
        "consensus_matrix": consensus_data
    }