from __future__ import annotations

import datetime
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

import yaml

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

if TYPE_CHECKING:
    from ocr_pipeline.citations import Citation
    from ocr_pipeline.detectors.figures import DetectedFigure
    from ocr_pipeline.detectors.tables import DetectedTable
    from ocr_pipeline.postprocess.entities import EntityResult

logger = get_logger(__name__)


@dataclass
class OutputMetadata:
    source_path: str
    source_hash: str
    timestamp: str
    ocr_engine: str
    ocr_confidence_mean: float
    language: str
    detected_entities: list[str]
    has_figures: bool
    has_tables: bool
    citations_found: list[str]
    chunk_index: int
    total_chunks: int


class MarkdownWriter:
    def __init__(self) -> None:
        self.base_dir = Path(settings.output.base_directory).expanduser().resolve()
        self.organize_by = settings.output.organize_by
        self.frontmatter_fields = settings.output.frontmatter.fields
        self._run_dir: Path | None = None
        self._written_chunks: list[tuple[dict[str, Any], str]] = []

    def start_run(self, timestamp: str, run_number: int) -> Path:
        """Allocate one shared output directory for the entire pipeline run.

        The directory itself is created lazily on the first write, so a run
        where every image is skipped does not leave an empty run_NNN behind.
        """
        self._written_chunks = []
        dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
        if self.organize_by == "date_run":
            self._run_dir = self.base_dir / dt.strftime("%Y-%m-%d") / f"run_{run_number:03d}"
        elif self.organize_by == "flat":
            self._run_dir = self.base_dir
        else:
            self._run_dir = None
        if self._run_dir is not None:
            return self._run_dir
        return self.base_dir

    def _get_output_dir(self, metadata: dict[str, Any]) -> Path:
        if self.organize_by == "source_dir":
            source_parent = Path(str(metadata.get("source_path", "unknown"))).parent
            safe_parent = "_".join(
                part for part in source_parent.parts if part not in (source_parent.anchor, "/")
            )
            output_dir = self.base_dir / "by_source" / (safe_parent[:100] or "unknown")
            output_dir.mkdir(parents=True, exist_ok=True)
            return output_dir
        if self._run_dir is not None:
            self._run_dir.mkdir(parents=True, exist_ok=True)
            return self._run_dir
        timestamp = metadata.get("timestamp")
        if not isinstance(timestamp, str) or not timestamp:
            timestamp = datetime.datetime.now(datetime.UTC).isoformat()
        return self.start_run(timestamp, 1)

    def write_chunk(
        self,
        *,
        content: str,
        metadata: dict[str, Any],
        figures: list[DetectedFigure] | None = None,
        tables: list[DetectedTable] | None = None,
        entities: EntityResult | None = None,
        citations: list[Citation] | None = None,
    ) -> Path:
        output_dir = self._get_output_dir(metadata)
        source_path = Path(str(metadata.get("source_path", "unknown")))
        safe_name = (
            "".join(c for c in source_path.stem if c.isalnum() or c in "-_")[:80] or "unknown"
        )
        source_hash = str(metadata.get("source_hash", ""))[:12]
        chunk_idx = int(metadata.get("chunk_index", 0))
        total_chunks = int(metadata.get("total_chunks", 1))
        suffix = f"_{source_hash}" if source_hash else ""
        filename = (
            f"{safe_name}{suffix}_chunk_{chunk_idx:03d}.md"
            if total_chunks > 1
            else f"{safe_name}{suffix}.md"
        )
        output_path = output_dir / filename

        frontmatter = self._build_frontmatter(metadata)
        body = self._build_body(
            content, source_path.name, figures or [], tables or [], entities, citations or []
        )
        self._written_chunks.append((dict(metadata), content))
        with output_path.open("w", encoding="utf-8") as handle:
            handle.write("---\n")
            yaml.safe_dump(frontmatter, handle, sort_keys=False, allow_unicode=True)
            handle.write("---\n\n")
            handle.write(body)

        logger.debug("markdown_written", path=str(output_path))
        return output_path

    def write_consolidated(self) -> Path | None:
        """Write one readable document while retaining chunk files for RAG."""
        if not self._written_chunks:
            return None
        output_dir = self._run_dir or self.base_dir
        output_dir.mkdir(parents=True, exist_ok=True)
        grouped: dict[str, list[tuple[dict[str, Any], str]]] = {}
        for metadata, content in self._written_chunks:
            grouped.setdefault(str(metadata.get("source_path", "unknown")), []).append(
                (metadata, content)
            )

        timestamp = datetime.datetime.now(datetime.UTC).isoformat()
        frontmatter = {
            "document_type": "consolidated_ocr",
            "generated_at": timestamp,
            "source_count": len(grouped),
            "chunk_count": len(self._written_chunks),
            "sources": [str(path) for path in grouped],
        }
        parts = ["# Consolidated OCR Output", ""]
        for source_path, chunks in sorted(grouped.items()):
            parts.extend([f"## Source: {Path(source_path).name}", ""])
            for _metadata, content in sorted(
                chunks, key=lambda item: int(item[0].get("chunk_index", 0))
            ):
                parts.extend([content.strip(), ""])
        output_path = output_dir / settings.output.consolidated_filename
        with output_path.open("w", encoding="utf-8") as handle:
            handle.write("---\n")
            yaml.safe_dump(frontmatter, handle, sort_keys=False, allow_unicode=True)
            handle.write("---\n\n")
            handle.write("\n".join(parts).rstrip() + "\n")
        logger.info(
            "consolidated_markdown_written",
            path=str(output_path),
            sources=len(grouped),
            chunks=len(self._written_chunks),
        )
        return output_path

    def _build_frontmatter(self, metadata: dict[str, Any]) -> dict[str, Any]:
        return {field: metadata[field] for field in self.frontmatter_fields if field in metadata}

    def _build_body(
        self,
        content: str,
        source_name: str,
        figures: list[DetectedFigure],
        tables: list[DetectedTable],
        entities: EntityResult | None,
        citations: list[Citation],
    ) -> str:
        parts = [f"# Screenshot: {source_name}", "", "## Extracted Text", "", content.strip(), ""]
        if figures:
            parts.extend(["## Figures Detected", ""])
            for index, figure in enumerate(figures, 1):
                parts.extend(
                    [
                        f"### Figure {index}",
                        f"- **Bounding Box**: {figure.bbox}",
                        f"- **Confidence**: {figure.confidence:.2f}",
                    ]
                )
                if figure.caption:
                    parts.append(f"- **Caption**: {figure.caption}")
                parts.append("")
        if tables:
            parts.extend(["## Tables Detected", ""])
            for index, table in enumerate(tables, 1):
                parts.extend(
                    [
                        f"### Table {index}",
                        f"- **Bounding Box**: {table.bbox}",
                        f"- **Confidence**: {table.confidence:.2f}",
                    ]
                )
                if table.markdown:
                    parts.extend(["", "#### Table Content (Markdown)", "", table.markdown])
                parts.append("")
        if entities and entities.entities:
            parts.extend(
                [
                    "## Detected Entities",
                    "",
                    "| Entity | Type | Context |",
                    "|--------|------|---------|",
                ]
            )
            for entity in entities.entities[:50]:
                context = entity.context[:80].replace("|", "\\|")
                if len(entity.context) > 80:
                    context += "..."
                parts.append(f"| {entity.text} | {entity.label} | {context} |")
            parts.append("")
        if citations:
            parts.extend(
                ["## Citations Found", "", "| Type | Identifier |", "|------|------------|"]
            )
            for citation in citations:
                parts.append(f"| {citation.type.upper()} | `{citation.identifier}` |")
            parts.append("")
        return "\n".join(parts)
