from __future__ import annotations

import datetime
from pathlib import Path
from typing import Any

import yaml

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

logger = get_logger(__name__)


class MarkdownWriter:
    def __init__(self):
        self.base_dir = Path(settings.output.base_directory).expanduser().resolve()
        self.organize_by = settings.output.organize_by
        self.frontmatter_fields = settings.output.frontmatter.fields

    def _get_output_dir(self, timestamp: str) -> Path:
        dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00"))

        if self.organize_by == "date_run":
            date_dir = self.base_dir / dt.strftime("%Y-%m-%d")
            run_dirs = sorted(date_dir.glob("run_*"))
            run_num = len(run_dirs) + 1
            return date_dir / f"run_{run_num:03d}"
        elif self.organize_by == "source_dir":
            return self.base_dir / "by_source"
        else:
            return self.base_dir

    def write_chunk(
        self,
        content: str,
        metadata: dict[str, Any],
        figures: list[Any] = None,
        tables: list[Any] = None,
        entities: Any = None,
        citations: list[Any] = None,
    ) -> Path:
        output_dir = self._get_output_dir(metadata.get("timestamp", ""))
        output_dir.mkdir(parents=True, exist_ok=True)

        source_path = Path(metadata.get("source_path", "unknown"))
        safe_name = "".join(c for c in source_path.stem if c.isalnum() or c in "-_")[:80]
        chunk_idx = metadata.get("chunk_index", 0)
        total_chunks = metadata.get("total_chunks", 1)

        if total_chunks > 1:
            filename = f"{safe_name}_chunk_{chunk_idx:03d}.md"
        else:
            filename = f"{safe_name}.md"

        output_path = output_dir / filename

        frontmatter = self._build_frontmatter(metadata)
        body = self._build_body(content, figures, tables, entities, citations)

        with open(output_path, "w") as f:
            f.write("---\n")
            yaml.dump(frontmatter, f, sort_keys=False, allow_unicode=True)
            f.write("---\n\n")
            f.write(body)

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

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

    def _build_body(
        self,
        content: str,
        figures: list[Any] = None,
        tables: list[Any] = None,
        entities: Any = None,
        citations: list[Any] = None,
    ) -> str:
        parts = []

        source_name = Path(content.split("\n")[0]).name if content else "Unknown"
        parts.append(f"# Screenshot: {source_name}")
        parts.append("")

        if content:
            parts.append("## Extracted Text")
            parts.append("")
            parts.append(content.strip())
            parts.append("")

        if figures:
            parts.append("## Figures Detected")
            parts.append("")
            for i, fig in enumerate(figures):
                parts.append(f"### Figure {i + 1}")
                parts.append(f"- **Bounding Box**: {fig.bbox}")
                parts.append(f"- **Confidence**: {fig.confidence:.2f}")
                if fig.caption:
                    parts.append(f"- **Caption**: {fig.caption}")
                parts.append("")

        if tables:
            parts.append("## Tables Detected")
            parts.append("")
            for i, table in enumerate(tables):
                parts.append(f"### Table {i + 1}")
                parts.append(f"- **Bounding Box**: {table.bbox}")
                parts.append(f"- **Confidence**: {table.confidence:.2f}")
                if table.markdown:
                    parts.append("")
                    parts.append("#### Table Content (Markdown)")
                    parts.append("")
                    parts.append(table.markdown)
                parts.append("")

        if entities and entities.entities:
            parts.append("## Detected Entities")
            parts.append("")
            parts.append("| Entity | Type | Context |")
            parts.append("|--------|------|---------|")
            for ent in entities.entities[:50]:
                context = ent.context[:80].replace("|", "\\|") + (
                    "..." if len(ent.context) > 80 else ""
                )
                parts.append(f"| {ent.text} | {ent.label} | {context} |")
            parts.append("")

        if citations:
            parts.append("## Citations Found")
            parts.append("")
            parts.append("| Type | Identifier |")
            parts.append("|------|------------|")
            for cit in citations:
                parts.append(f"| {cit.type.upper()} | `{cit.identifier}` |")
            parts.append("")

        return "\n".join(parts)