from __future__ import annotations

import re
from dataclasses import dataclass

import spacy

from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger

logger = get_logger(__name__)


@dataclass
class Entity:
    text: str
    label: str
    start: int
    end: int
    confidence: float
    context: str = ""


@dataclass
class EntityResult:
    entities: list[Entity]
    text: str
    total_entities: int
    entities_by_type: dict[str, int]


class ScientificEntityRecognizer:
    def __init__(self):
        self._nlp = None
        self._initialized = False

    def _init(self):
        if self._initialized:
            return
        try:

            self._nlp = spacy.load(settings.entities.model)
            self._nlp.max_length = 2_000_000
            self._initialized = True
            logger.info("scispacy_model_loaded", model=settings.entities.model)
        except Exception as e:
            logger.warning("scispacy_load_failed", error=str(e), model=settings.entities.model)
            self._fallback_init()
            self._initialized = True

    def _fallback_init(self):
        try:
            self._nlp = spacy.blank("en")
            self._nlp.add_pipe("sentencizer")
            logger.info("using_fallback_blank_model")
        except Exception:
            self._nlp = None

    def extract(self, text: str) -> EntityResult:
        if not settings.entities.enabled:
            return EntityResult(entities=[], text=text, total_entities=0, entities_by_type={})

        self._init()
        if self._nlp is None:
            return EntityResult(entities=[], text=text, total_entities=0, entities_by_type={})

        try:
            doc = self._nlp(text)
            entities = []

            for ent in doc.ents:
                if ent.label_ in settings.entities.types:
                    context_start = max(0, ent.start_char - 100)
                    context_end = min(len(text), ent.end_char + 100)
                    context = text[context_start:context_end].strip()

                    entities.append(
                        Entity(
                            text=ent.text,
                            label=ent.label_,
                            start=ent.start_char,
                            end=ent.end_char,
                            confidence=0.9,
                            context=context,
                        )
                    )

            entities = (
                self._merge_entities(entities, text) if settings.entities.merge_entities else entities
            )
            entities = self._deduplicate_entities(entities)

            by_type = {}
            for e in entities:
                by_type[e.label] = by_type.get(e.label, 0) + 1

            return EntityResult(
                entities=entities,
                text=text,
                total_entities=len(entities),
                entities_by_type=by_type,
            )
        except Exception as e:
            logger.warning("entity_extraction_failed", error=str(e))
            return EntityResult(entities=[], text=text, total_entities=0, entities_by_type={})

    def _merge_entities(self, entities: list[Entity], text: str) -> list[Entity]:
        if not entities:
            return entities

        entities.sort(key=lambda e: (e.start, -e.end))
        merged = [entities[0]]

        for ent in entities[1:]:
            last = merged[-1]
            if ent.start <= last.end + 5 and ent.label == last.label:
                merged[-1] = Entity(
                    text=last.text + text[last.end : ent.start] + ent.text,
                    label=last.label,
                    start=last.start,
                    end=ent.end,
                    confidence=max(last.confidence, ent.confidence),
                    context=last.context,
                )
            else:
                merged.append(ent)

        return merged

    def _deduplicate_entities(self, entities: list[Entity]) -> list[Entity]:
        seen = set()
        unique = []
        for ent in entities:
            key = (ent.text.lower(), ent.label)
            if key not in seen:
                seen.add(key)
                unique.append(ent)
        return unique


class RegexEntityRecognizer:
    PATTERNS = {
        "GENE": [
            r"\b(?:[A-Z]{2,5}\d*[A-Z]?|[A-Z]{3,10})\b",
            r"\b(?:BRCA[12]|TP53|EGFR|KRAS|MYC|PIK3CA|PTEN|APC|MLH1|MSH2|MSH6|PMS2|ATM|CHEK2|PALB2|RAD51|BARD1|BRIP1|CDH1|STK11|VHL|WT1|NF1|NF2|RB1|MEN1|RET|MET|ALK|ROS1|NTRK[123]|BRAF|NRAS|HRAS|IDH[12]|TERT|FGFR[1234]|KIT|PDGFRA|FLT3|JAK2|MPL|CALR|ASXL1|DNMT3A|TET2|SRSF2|U2AF1|SF3B1|ZRSR2|CBL|CSF3R|RUNX1|CEBPA|NPM1|FLT3|DNMT3A|IDH[12]|TET2|ASXL1|SRSF2|U2AF1|SF3B1|ZRSR2|CBL|CSF3R|RUNX1|CEBPA|NPM1)\b",
        ],
        "PROTEIN": [
            r"\b(?:p53|p21|p16|CDK[1246]|cyclin\s+[ADEB]|Rb|E2F|MDM2|BAX|BCL[2XL]|CASP[389]|PARP|H2AX|ATM|ATR|CHK[12]|BRCA[12]|RAD51|FANCD2|FANCI|PCNA|RPA|MSH[26]|MLH1|PMS2|EXO1|MRE11|RAD50|NBN|CTIP|BRCA1|BRCA2|PALB2|RAD51C|RAD51D|BRIP1|BARD1)\b",
        ],
        "CHEMICAL": [
            r"\b(?:DMSO|PBS|EDTA|Tris|HEPES|SDS|DTT|BME|NaCl|KCl|MgCl2|CaCl2|NaOH|HCl|H2SO4|HNO3|EtOH|MeOH|IPA|DMSO|DMF|DMA|THF|DCM|CHCl3|CH2Cl2|EtOAc|hexane|pentane|acetone|acetonitrile|water|H2O|buffer|media|serum|FBS|BSA|pen/strep|penicillin|streptomycin|trypsin|EDTA|collagenase|dispase|accutase)\b",
            r"\b(?:[A-Z][a-z]?\d*(?:\([^)]+\))?(?:\s*[+\-]\s*[A-Z][a-z]?\d*(?:\([^)]+\))?)*)\b",
        ],
        "CONCENTRATION": [
            r"\b\d+(?:\.\d+)?\s*(?:[µumMpnc]?[Mm]|[µumMpnc]?[Mm]/[Ll]|[µumMpnc]?[Mm]\s*[Ll]?|[µumMpnc]?[gG]\s*/\s*[Ll]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dL]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dD][lL]|[µumMpnc]?[gG]\s*/\s*100\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[kK][gG])\b",
        ],
    }

    @classmethod
    def extract(cls, text: str) -> list[Entity]:
        entities = []
        for label, patterns in cls.PATTERNS.items():
            for pattern in patterns:
                for match in re.finditer(pattern, text, re.IGNORECASE):
                    context_start = max(0, match.start() - 50)
                    context_end = min(len(text), match.end() + 50)
                    entities.append(
                        Entity(
                            text=match.group(),
                            label=label,
                            start=match.start(),
                            end=match.end(),
                            confidence=0.7,
                            context=text[context_start:context_end].strip(),
                        )
                    )
        return entities


def extract_entities(text: str) -> EntityResult:
    recognizer = ScientificEntityRecognizer()
    result = recognizer.extract(text)

    regex_entities = RegexEntityRecognizer.extract(text)
    all_entities = result.entities + regex_entities
    all_entities = _deduplicate(all_entities)

    by_type = {}
    for e in all_entities:
        by_type[e.label] = by_type.get(e.label, 0) + 1

    return EntityResult(
        entities=all_entities,
        text=text,
        total_entities=len(all_entities),
        entities_by_type=by_type,
    )


def _deduplicate(entities: list[Entity]) -> list[Entity]:
    seen = set()
    unique = []
    for ent in sorted(entities, key=lambda e: (e.start, -e.end)):
        key = (ent.text.lower(), ent.label)
        if key not in seen:
            seen.add(key)
            unique.append(ent)
    return unique


def format_entities_for_markdown(result: EntityResult) -> str:
    if not result.entities:
        return "No entities detected."

    lines = [
        "## Detected Entities",
        "",
        "| Entity | Type | Count | Contexts |",
        "|--------|------|-------|----------|",
    ]

    for label in sorted(result.entities_by_type.keys()):
        ents = [e for e in result.entities if e.label == label]
        contexts = {e.context[:80] for e in ents[:3]}
        lines.append(f"| {ents[0].text} | {label} | {len(ents)} | {'; '.join(contexts)} |")

    return "\n".join(lines)
