# utils/verifier.py
from lettucedetect.models.inference import HallucinationDetector
import logging
from typing import List, Dict, Optional

logger = logging.getLogger(__name__)

class DischargeVerifier:
    def __init__(self):
        """Initialize the hallucination detector."""
        try:
            self.detector = HallucinationDetector(
                method="transformer",
                model_path="KRLabsOrg/lettucedect-base-modernbert-en-v1",
            )
            logger.info("Hallucination detector initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize hallucination detector: {str(e)}")
            raise

    def create_interactive_text(self, text: str, spans: List[Dict[str, int | float]]) -> str:
        """Create interactive HTML with highlighting and hover effects."""
        html_text = text

        for span in sorted(spans, key=lambda x: x["start"], reverse=True):
            span_text = text[span["start"]:span["end"]]
            highlighted_span = (
                f'<span class="hallucination" title="Confidence: {span["confidence"]:.3f}">{span_text}</span>'
            )
            html_text = (
                html_text[:span["start"]] + highlighted_span + html_text[span["end"]:]
            )

        return f"""
        <style>
            .container {{
                font-family: Arial, sans-serif;
                font-size: 16px;
                line-height: 1.6;
                padding: 20px;
            }}
            .hallucination {{
                background-color: rgba(255, 99, 71, 0.3);
                padding: 2px;
                border-radius: 3px;
                cursor: help;
            }}
            .hallucination:hover {{
                background-color: rgba(255, 99, 71, 0.5);
            }}
        </style>
        <div class="container">{html_text}</div>
        """

    def verify_discharge_summary(
        self, context: str, question: str, answer: str
    ) -> Optional[str]:
        """Verify the discharge summary for hallucinations and return highlighted HTML."""
        try:
            predictions = self.detector.predict(
                context=[context],
                question=question,
                answer=answer,
                output_format="spans"
            )
            logger.debug(f"Hallucination predictions: {predictions}")
            return self.create_interactive_text(answer, predictions)
        except Exception as e:
            logger.error(f"Error verifying discharge summary: {str(e)}")
            return None