import os
import time
from typing import List, Dict, Any
import numpy as np
import pandas as pd
import mlflow
from .ocr_space_engine import OCRSpaceEngine

def compute_cer(reference: str, hypothesis: str) -> float:
    """Computes Character Error Rate (CER) via edit distance."""
    r = reference.lower().strip()
    h = hypothesis.lower().strip()
    if not r:
        return 0.0 if not h else 1.0

    n, m = len(r), len(h)
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(n + 1):
        dp[i][0] = i
    for j in range(m + 1):
        dp[0][j] = j

    for i in range(1, n + 1):
        for j in range(1, m + 1):
            if r[i - 1] == h[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])

    dist = dp[n][m]
    return float(dist / max(len(r), 1))

def run_ocr_comparison(
    sample_records: List[Any],
    cache_dir: str = "data/processed/ocr_cache",
    tracking_uri: str = "sqlite:///mlflow.db"
) -> pd.DataFrame:
    """
    Runs comparison between OCR engines and logs results to MLflow under 'ocr_comparison'.
    """
    mlflow.set_tracking_uri(tracking_uri)
    exp = mlflow.set_experiment("ocr_comparison")

    engines = {
        "ocr_space_engine2": OCRSpaceEngine(ocr_engine="2", cache_dir=cache_dir),
        "ocr_space_engine1": OCRSpaceEngine(ocr_engine="1", cache_dir=cache_dir)
    }

    comparison_rows = []

    with mlflow.start_run(run_name="ocr_engine_benchmark", experiment_id=exp.experiment_id):
        for name, engine in engines.items():
            runtimes = []
            confidences = []
            cers = []

            for rec in sample_records:
                img_path = rec.image_path if hasattr(rec, "image_path") else rec.get("image_path")
                doc_id = rec.doc_id if hasattr(rec, "doc_id") else rec.get("doc_id")
                gt_words = rec.ocr_words if hasattr(rec, "ocr_words") else rec.get("ocr_words", [])
                gt_text = " ".join([w.text if hasattr(w, "text") else w.get("text", "") for w in gt_words])

                t0 = time.time()
                words, _ = engine.extract_with_cache(img_path, doc_id=doc_id)
                dur = time.time() - t0
                runtimes.append(dur)

                if words:
                    confs = [w.conf for w in words]
                    confidences.append(float(np.mean(confs)))
                    pred_text = " ".join([w.text for w in words])
                    if gt_text:
                        cers.append(compute_cer(gt_text, pred_text))
                    else:
                        cers.append(0.0)
                else:
                    confidences.append(0.0)
                    cers.append(1.0)

            avg_runtime = float(np.mean(runtimes)) if runtimes else 0.0
            avg_conf = float(np.mean(confidences)) if confidences else 0.0
            avg_cer = float(np.mean(cers)) if cers else 0.0

            mlflow.log_metric(f"{name}_avg_runtime_sec", avg_runtime)
            mlflow.log_metric(f"{name}_avg_confidence", avg_conf)
            mlflow.log_metric(f"{name}_cer", avg_cer)

            comparison_rows.append({
                "engine": name,
                "avg_runtime_sec": round(avg_runtime, 4),
                "avg_confidence": round(avg_conf, 4),
                "cer": round(avg_cer, 4)
            })

        df_res = pd.DataFrame(comparison_rows)
        csv_path = "reports/tables/ocr_engine_comparison.csv"
        os.makedirs(os.path.dirname(csv_path), exist_ok=True)
        df_res.to_csv(csv_path, index=False)
        mlflow.log_artifact(csv_path)

    return df_res

if __name__ == "__main__":
    from ..data.schema import DocumentRecord, OCRWord
    dummy_records = [
        DocumentRecord(
            doc_id="sample_0",
            source_dataset="test",
            doc_class="invoice",
            image_path="",
            ocr_words=[OCRWord(text="INVOICE TOTAL 100", bbox=[0, 0, 10, 10], conf=1.0)],
            fields={}
        )
    ]
    df = run_ocr_comparison(dummy_records)
    print(df)