from __future__ import annotations

import re
from dataclasses import dataclass

import httpx

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

logger = get_logger(__name__)


@dataclass
class Citation:
    raw: str
    type: str
    identifier: str
    confidence: float
    context: str = ""


class RegexCitationExtractor:
    DOI_PATTERN = re.compile(
        r"\b(?:10\.\d{4,9}/[-._;()/:A-Z0-9]+)\b",
        re.IGNORECASE,
    )
    PMID_PATTERN = re.compile(
        r"\b(?:PMID|PubMed\s*ID|PubMed)\s*[:\-]?\s*(\d{8})\b",
        re.IGNORECASE,
    )
    ARXIV_PATTERN = re.compile(
        r"\b(?:arXiv|arxiv)\s*[:\-]?\s*(\d{4}\.\d{4,5}(?:v\d+)?)\b",
        re.IGNORECASE,
    )
    ISBN_PATTERN = re.compile(
        r"\b(?:ISBN[-\s]*(?:13|10)?[-\s]*)?(?:97[89][-\s]*)?\d{1,5}[-\s]*\d{1,7}[-\s]*\d{1,7}[-\s]*[\dX]\b",
        re.IGNORECASE,
    )

    @classmethod
    def extract(cls, text: str) -> list[Citation]:
        citations = []

        for match in cls.DOI_PATTERN.finditer(text):
            ctx_start = max(0, match.start() - 80)
            ctx_end = min(len(text), match.end() + 80)
            citations.append(
                Citation(
                    raw=match.group(),
                    type="doi",
                    identifier=match.group().lower(),
                    confidence=0.95,
                    context=text[ctx_start:ctx_end].strip(),
                )
            )

        for match in cls.PMID_PATTERN.finditer(text):
            ctx_start = max(0, match.start() - 80)
            ctx_end = min(len(text), match.end() + 80)
            citations.append(
                Citation(
                    raw=match.group(),
                    type="pmid",
                    identifier=f"PMID:{match.group(1)}",
                    confidence=0.9,
                    context=text[ctx_start:ctx_end].strip(),
                )
            )

        for match in cls.ARXIV_PATTERN.finditer(text):
            ctx_start = max(0, match.start() - 80)
            ctx_end = min(len(text), match.end() + 80)
            citations.append(
                Citation(
                    raw=match.group(),
                    type="arxiv",
                    identifier=f"arXiv:{match.group(1)}",
                    confidence=0.9,
                    context=text[ctx_start:ctx_end].strip(),
                )
            )

        for match in cls.ISBN_PATTERN.finditer(text):
            ctx_start = max(0, match.start() - 80)
            ctx_end = min(len(text), match.end() + 80)
            isbn = re.sub(r"[-\s]", "", match.group())
            citations.append(
                Citation(
                    raw=match.group(),
                    type="isbn",
                    identifier=f"ISBN:{isbn}",
                    confidence=0.7,
                    context=text[ctx_start:ctx_end].strip(),
                )
            )

        return citations


class GrobidClient:
    def __init__(self):
        self.url = settings.citations.grobid_url
        self.timeout = settings.citations.grobid_timeout

    def process_text(self, text: str) -> list[Citation]:
        if not settings.citations.grobid_enabled:
            return []

        try:
            with httpx.Client(timeout=self.timeout) as client:
                response = client.post(
                    f"{self.url}/api/processCitationText",
                    data={"text": text},
                )
                response.raise_for_status()
                return self._parse_grobid_xml(response.text)
        except Exception as e:
            logger.warning("grobid_request_failed", error=str(e))
            return []

    def _parse_grobid_xml(self, xml_text: str) -> list[Citation]:
        try:
            import xml.etree.ElementTree as ET

            root = ET.fromstring(xml_text)
            citations = []

            ns = {"tei": "http://www.tei-c.org/ns/1.0"}
            for cit in root.findall(".//tei:biblStruct", ns):
                doi_elem = cit.find(".//tei:idno[@type='DOI']", ns)
                pmid_elem = cit.find(".//tei:idno[@type='PMID']", ns)

                if doi_elem is not None and doi_elem.text:
                    citations.append(
                        Citation(
                            raw=doi_elem.text,
                            type="doi",
                            identifier=doi_elem.text.lower(),
                            confidence=0.99,
                        )
                    )
                if pmid_elem is not None and pmid_elem.text:
                    citations.append(
                        Citation(
                            raw=pmid_elem.text,
                            type="pmid",
                            identifier=f"PMID:{pmid_elem.text}",
                            confidence=0.99,
                        )
                    )

            return citations
        except Exception as e:
            logger.warning("grobid_xml_parse_failed", error=str(e))
            return []


def extract_citations(text: str) -> list[Citation]:
    citations = []

    if settings.citations.regex_enabled:
        citations.extend(RegexCitationExtractor.extract(text))

    if settings.citations.grobid_enabled:
        client = GrobidClient()
        citations.extend(client.process_text(text))

    seen = set()
    unique = []
    for cit in citations:
        key = (cit.type, cit.identifier)
        if key not in seen:
            seen.add(key)
            unique.append(cit)

    return unique


def format_citations_for_markdown(citations: list[Citation]) -> str:
    if not citations:
        return "No citations found."

    lines = [
        "## Citations",
        "",
        "| Type | Identifier | Context |",
        "|------|------------|---------|",
    ]
    for cit in citations:
        context = (
            cit.context[:80].replace("|", "\\|") + "..." if len(cit.context) > 80 else cit.context
        )
        lines.append(f"| {cit.type.upper()} | `{cit.identifier}` | {context} |")

    return "\n".join(lines)
