"""Pipeline Engine — orchestrates the full Meeting-to-Ontology flow.

Usage:
    engine = await PipelineEngine.create(config)
    result = await engine.ingest("path/to/meeting.md")
    result = await engine.rerun("DOC-abc123", model="gemini-2.5-pro")
"""

from __future__ import annotations

import json
import time
from datetime import date
from pathlib import Path
from typing import Any

from ontology_engine.core.config import OntologyConfig
from ontology_engine.core.errors import ExtractionError
from ontology_engine.core.types import (
    Entity,
    EntityType,
    ExtractionResult,
    Link,
    LinkType,
    Provenance,
    ValidationResult,
)
from ontology_engine.llm.base import LLMProvider
from ontology_engine.pipeline.extractor import StructuredExtractor
from ontology_engine.pipeline.preprocessor import MeetingPreprocessor, ProcessedMeeting
from ontology_engine.pipeline.validator import ExtractionValidator
from ontology_engine.storage.bronze import BronzeRepository


class PipelineEngine:
    """End-to-end pipeline: file → bronze → preprocess → extract → validate → store.

    Accepts an optional domain_schema to drive schema-aware extraction,
    validation, and provenance tracking.
    """

    def __init__(
        self,
        llm: LLMProvider,
        repo: Any | None,
        config: OntologyConfig,
        bronze: BronzeRepository | None = None,
        domain_schema: Any | None = None,
    ):
        self.llm = llm
        self.repo = repo
        self.bronze = bronze
        self.config = config
        self._domain_schema = domain_schema
        self.preprocessor = MeetingPreprocessor(llm, config)
        self.extractor = StructuredExtractor(llm, config, domain_schema=domain_schema)
        self.validator = ExtractionValidator(config, domain_schema=domain_schema)

    @classmethod
    async def create(
        cls,
        config: OntologyConfig,
        db_url: str | None = None,
        domain_schema: Any | None = None,
    ) -> PipelineEngine:
        """Factory: create engine with LLM provider and optional DB."""
        from ontology_engine.storage.repository import OntologyRepository

        llm = _create_llm(config)

        repo = None
        bronze = None
        url = db_url or config.database.url
        if url:
            repo = await OntologyRepository.create(url, config.database.db_schema)
            bronze = BronzeRepository(repo._pool, config.database.db_schema)

        return cls(llm, repo, config, bronze=bronze, domain_schema=domain_schema)

    # -----------------------------------------------------------------
    # Schema resolution helpers
    # -----------------------------------------------------------------

    def _resolve_domain_schema(self, schema: Any | None) -> Any | None:
        """Resolve *schema* to a DomainSchema object.

        Accepts None, a DomainSchema object, a dict, or a domain name string.
        """
        if schema is None:
            return None
        if hasattr(schema, "entity_types") and hasattr(schema, "domain"):
            return schema
        if isinstance(schema, dict):
            from ontology_engine.core.schema_registry import DomainSchema
            return DomainSchema.from_dict(schema)
        if isinstance(schema, str):
            return self._load_schema_by_name(schema)
        return None

    @staticmethod
    def _load_schema_by_name(name: str) -> Any:
        """Load a domain schema YAML by domain name."""
        from ontology_engine.core.schema_registry import DomainSchema

        candidates = [
            Path(__file__).resolve().parent.parent.parent.parent / "domain_schemas",
            Path.cwd() / "domain_schemas",
        ]
        for schema_dir in candidates:
            for ext in (".yaml", ".yml"):
                p = schema_dir / f"{name}{ext}"
                if p.exists():
                    return DomainSchema.from_yaml(p)
        raise ExtractionError(
            f"Domain schema '{name}' not found in: {[str(c) for c in candidates]}"
        )

    # -----------------------------------------------------------------
    # Ingest
    # -----------------------------------------------------------------

    async def ingest(
        self,
        file_path: str,
        *,
        meeting_date: date | None = None,
        participants: list[str] | None = None,
        source_type: str = "meeting_transcript",
        content_format: str = "text",
        schema: Any | None = None,
    ) -> IngestResult:
        """Ingest a single meeting transcript file.

        Args:
            schema: Optional per-call domain schema override.  Accepts a
                DomainSchema object, a dict, or a schema name string
                (e.g. ``"edtech"``).  When *None*, the engine-level schema
                (set at construction time) is used.

        Flow: file → Bronze (dedup) → preprocess → extract → validate → store.
        """
        t0 = time.monotonic()

        path = Path(file_path)
        if not path.exists():
            raise ExtractionError(f"File not found: {file_path}")
        raw_text = path.read_text(encoding="utf-8")

        if meeting_date is None:
            meeting_date = self._parse_date_from_filename(path.name)

        # Resolve per-call schema override → fall back to engine-level schema
        effective_schema = self._resolve_domain_schema(schema) or self._domain_schema
        if effective_schema is not None and effective_schema is not self._domain_schema:
            extractor = StructuredExtractor(
                self.llm, self.config, domain_schema=effective_schema,
            )
            validator = ExtractionValidator(
                self.config, domain_schema=effective_schema,
            )
        else:
            extractor = self.extractor
            validator = self.validator

        # Stage 0: Bronze Layer — immutable raw document storage (dedup)
        bronze_doc_id: str | None = None
        if self.bronze:
            bronze_doc_id, is_new = await self.bronze.ingest(
                content=raw_text,
                source_type=source_type,
                source_uri=str(path),
                content_format=content_format,
                metadata={
                    "meeting_date": meeting_date.isoformat() if meeting_date else None,
                    "participants": participants or [],
                    "schema": getattr(effective_schema, "domain", None),
                },
                ingested_by="pipeline",
            )
            if not is_new:
                elapsed_ms = int((time.monotonic() - t0) * 1000)
                return IngestResult(
                    file=str(path),
                    meeting_date=meeting_date,
                    bronze_doc_id=bronze_doc_id,
                    skipped_duplicate=True,
                    processing_time_ms=elapsed_ms,
                )

        # Stage 1: Preprocessing
        processed = await self.preprocessor.process(
            raw_text,
            meeting_date=meeting_date,
            participants=participants,
        )
        processed.metadata["source_file"] = str(path)

        # Stage 2: Extraction (4-pass)
        extraction = await extractor.extract(processed)
        extraction.source_file = str(path)

        # Stage 3: Validation
        validation = validator.validate(extraction)

        # Stage 4: Store (if DB connected and validation passed)
        stored_ids: dict[str, list[str]] = {
            "entities": [],
            "links": [],
            "provenance": [],
        }
        if self.repo and validation.is_valid and validation.extraction:
            stored_ids = await self._store_results(
                validation.extraction, str(path), bronze_doc_id=bronze_doc_id
            )

        elapsed_ms = int((time.monotonic() - t0) * 1000)

        return IngestResult(
            file=str(path),
            meeting_date=meeting_date,
            participants=processed.participants,
            extraction=validation.extraction or extraction,
            validation=validation,
            stored_ids=stored_ids,
            bronze_doc_id=bronze_doc_id,
            processing_time_ms=elapsed_ms,
        )

    async def rerun(
        self,
        bronze_doc_id: str,
        *,
        model: str | None = None,
        domain_schema: Any | None = None,
    ) -> IngestResult:
        """Re-run extraction on a Bronze document.

        Old Silver results are marked as archived (not deleted).
        Optionally use a different LLM model or domain schema.

        Args:
            bronze_doc_id: The Bronze document ID to re-extract.
            model: Optional override LLM model name.
            domain_schema: Optional override domain schema.

        Returns:
            IngestResult with fresh extraction results.
        """
        t0 = time.monotonic()

        if self.bronze is None:
            raise ExtractionError("Bronze repository not available — cannot rerun")

        doc = await self.bronze.get(bronze_doc_id)
        if doc is None:
            raise ExtractionError(f"Bronze document not found: {bronze_doc_id}")

        # Resolve LLM and schema for this rerun
        llm = self.llm if model is None else _create_llm_with_model(self.config, model)
        schema = domain_schema or self._domain_schema
        extractor = StructuredExtractor(llm, self.config, domain_schema=schema)
        validator = ExtractionValidator(self.config, domain_schema=schema)

        # Parse metadata
        meeting_date = None
        participants: list[str] = []
        if doc.metadata:
            md = doc.metadata.get("meeting_date")
            if md:
                try:
                    meeting_date = date.fromisoformat(md)
                except (ValueError, TypeError):
                    pass
            participants = doc.metadata.get("participants", [])

        # Re-run preprocessing + extraction
        processed = await self.preprocessor.process(
            doc.content,
            meeting_date=meeting_date,
            participants=participants if participants else None,
        )
        processed.metadata["source_file"] = doc.source_uri or ""
        processed.metadata["rerun_from"] = bronze_doc_id

        extraction = await extractor.extract(processed)
        extraction.source_file = doc.source_uri or ""

        validation = validator.validate(extraction)

        # Archive old results and store new ones
        stored_ids: dict[str, list[str]] = {
            "entities": [],
            "links": [],
            "provenance": [],
        }
        if self.repo and validation.is_valid and validation.extraction:
            await self._supersede_old_results(bronze_doc_id)
            stored_ids = await self._store_results(
                validation.extraction,
                doc.source_uri or "",
                bronze_doc_id=bronze_doc_id,
            )

        elapsed_ms = int((time.monotonic() - t0) * 1000)

        return IngestResult(
            file=doc.source_uri or "",
            meeting_date=meeting_date,
            participants=participants,
            extraction=validation.extraction or extraction,
            validation=validation,
            stored_ids=stored_ids,
            bronze_doc_id=bronze_doc_id,
            processing_time_ms=elapsed_ms,
        )

    async def _supersede_old_results(self, bronze_doc_id: str) -> None:
        """Mark old entities associated with a Bronze doc as archived."""
        if self.repo is None:
            return
        s = self.repo._schema
        sql = f"""
            UPDATE {s}.ont_entities
            SET status = 'archived', updated_at = NOW(), updated_by = 'pipeline_rerun'
            WHERE id IN (
                SELECT entity_id FROM {s}.ont_provenance
                WHERE source_document_id = $1 AND entity_id IS NOT NULL
            ) AND status = 'active'
        """
        async with self.repo._pool.acquire() as conn:
            await conn.execute(sql, bronze_doc_id)

    async def ingest_directory(
        self, directory: str, pattern: str = "*.md"
    ) -> list[IngestResult]:
        """Ingest all matching files in a directory."""
        dir_path = Path(directory)
        results: list[IngestResult] = []
        for f in sorted(dir_path.glob(pattern)):
            try:
                result = await self.ingest(str(f))
                results.append(result)
            except Exception as e:
                results.append(IngestResult(file=str(f), error=str(e)))
        return results

    # =========================================================================
    # Storage
    # =========================================================================

    async def _store_results(
        self,
        extraction: ExtractionResult,
        source_file: str,
        *,
        bronze_doc_id: str | None = None,
    ) -> dict[str, list[str]]:
        """Persist extraction results to PostgreSQL."""
        assert self.repo is not None

        ids: dict[str, list[str]] = {"entities": [], "links": [], "provenance": []}
        entity_name_to_id: dict[str, str] = {}

        # Store entities
        for ext_ent in extraction.entities:
            entity = Entity(
                entity_type=ext_ent.entity_type,
                name=ext_ent.name,
                properties=ext_ent.properties,
                aliases=ext_ent.aliases,
                confidence=ext_ent.confidence,
                created_by="pipeline",
            )

            # Check for existing entity (upsert logic) — only for enum types
            et_for_query = (
                ext_ent.entity_type
                if isinstance(ext_ent.entity_type, EntityType)
                else None
            )
            existing = (
                await self.repo.find_entity_by_name(ext_ent.name, et_for_query)
                if et_for_query
                else []
            )
            if existing:
                ent_id = existing[0].id
                assert ent_id is not None
                await self.repo.update_entity(
                    ent_id,
                    {
                        "properties": {**existing[0].properties, **ext_ent.properties},
                        "aliases": list(set(existing[0].aliases + ext_ent.aliases)),
                        "confidence": max(existing[0].confidence, ext_ent.confidence),
                    },
                    updated_by="pipeline",
                )
            else:
                ent_id = await self.repo.create_entity(entity)

            entity_name_to_id[ext_ent.name.lower()] = ent_id
            ids["entities"].append(ent_id)

            # Provenance
            prov_id = await self.repo.create_provenance(
                Provenance(
                    entity_id=ent_id,
                    source_document_id=bronze_doc_id,
                    source_type="meeting_transcript",
                    source_file=source_file,
                    source_meeting_date=extraction.meeting_date,
                    source_participants=extraction.participants,
                    source_segment=ext_ent.context,
                    extraction_model=extraction.extraction_model,
                    extraction_schema=extraction.extraction_schema or None,
                    extraction_pass="pass1",
                    created_by="pipeline",
                )
            )
            ids["provenance"].append(prov_id)

        # Store decisions as entities
        for dec in extraction.decisions:
            dec_entity = Entity(
                entity_type=EntityType.DECISION,
                name=dec.summary,
                properties={
                    "detail": dec.detail,
                    "decision_type": dec.decision_type.value,
                    "rationale": dec.rationale,
                    "conditions": dec.conditions,
                    "status": "active",
                },
                confidence=dec.confidence,
                created_by="pipeline",
            )
            dec_id = await self.repo.create_entity(dec_entity)
            entity_name_to_id[dec.summary.lower()] = dec_id
            ids["entities"].append(dec_id)

            # Link: Person → makes → Decision
            if dec.made_by:
                maker_id = entity_name_to_id.get(dec.made_by.lower())
                if maker_id:
                    link = Link(
                        link_type=LinkType.MAKES,
                        source_entity_id=maker_id,
                        target_entity_id=dec_id,
                        created_by="pipeline",
                    )
                    link_id = await self.repo.create_link(link)
                    ids["links"].append(link_id)

        # Store action items as entities
        for act in extraction.action_items:
            act_entity = Entity(
                entity_type=EntityType.ACTION_ITEM,
                name=act.task,
                properties={
                    "priority": act.priority,
                    "status": "pending",
                    "due_date": act.due_date.isoformat() if act.due_date else None,
                    "completion_criteria": act.completion_criteria,
                },
                confidence=act.confidence,
                created_by="pipeline",
            )
            act_id = await self.repo.create_entity(act_entity)
            ids["entities"].append(act_id)

            # Link: ActionItem → assigned_to → Person
            if act.owner:
                owner_id = entity_name_to_id.get(act.owner.lower())
                if owner_id:
                    link = Link(
                        link_type=LinkType.ASSIGNED_TO,
                        source_entity_id=act_id,
                        target_entity_id=owner_id,
                        created_by="pipeline",
                    )
                    link_id = await self.repo.create_link(link)
                    ids["links"].append(link_id)

        # Store extracted links
        for ext_link in extraction.links:
            src_id = entity_name_to_id.get(ext_link.source_name.lower())
            tgt_id = entity_name_to_id.get(ext_link.target_name.lower())
            if src_id and tgt_id:
                link = Link(
                    link_type=ext_link.link_type,
                    source_entity_id=src_id,
                    target_entity_id=tgt_id,
                    properties=ext_link.properties,
                    confidence=ext_link.confidence,
                    created_by="pipeline",
                )
                try:
                    link_id = await self.repo.create_link(link)
                    ids["links"].append(link_id)
                except Exception:
                    pass  # Duplicate link constraint

        return ids

    # =========================================================================
    # Helpers
    # =========================================================================

    @staticmethod
    def _parse_date_from_filename(filename: str) -> date | None:
        """Try to extract a date from filename like '20260208_meeting.md'."""
        import re

        m = re.search(r"(\d{4})(\d{2})(\d{2})", filename)
        if m:
            try:
                return date(int(m.group(1)), int(m.group(2)), int(m.group(3)))
            except ValueError:
                pass

        m = re.search(r"(\d{4})-(\d{2})-(\d{2})", filename)
        if m:
            try:
                return date(int(m.group(1)), int(m.group(2)), int(m.group(3)))
            except ValueError:
                pass

        return None

    async def close(self) -> None:
        """Clean up resources."""
        if self.repo:
            await self.repo.close()
        if self.bronze and not self.repo:
            await self.bronze.close()


class IngestResult:
    """Result of processing a single meeting file."""

    def __init__(
        self,
        file: str = "",
        meeting_date: date | None = None,
        participants: list[str] | None = None,
        extraction: ExtractionResult | None = None,
        validation: ValidationResult | None = None,
        stored_ids: dict[str, list[str]] | None = None,
        processing_time_ms: int = 0,
        error: str | None = None,
        bronze_doc_id: str | None = None,
        skipped_duplicate: bool = False,
    ):
        self.file = file
        self.meeting_date = meeting_date
        self.participants = participants or []
        self.extraction = extraction
        self.validation = validation
        self.stored_ids = stored_ids or {}
        self.processing_time_ms = processing_time_ms
        self.error = error
        self.bronze_doc_id = bronze_doc_id
        self.skipped_duplicate = skipped_duplicate

    @property
    def success(self) -> bool:
        return self.error is None

    def summary(self) -> dict[str, Any]:
        """Quick summary for logging."""
        s: dict[str, Any] = {
            "file": self.file,
            "success": self.success,
            "time_ms": self.processing_time_ms,
        }
        if self.bronze_doc_id:
            s["bronze_doc_id"] = self.bronze_doc_id
        if self.skipped_duplicate:
            s["skipped_duplicate"] = True
        if self.error:
            s["error"] = self.error
        if self.extraction:
            s["entities"] = len(self.extraction.entities)
            s["links"] = len(self.extraction.links)
            s["decisions"] = len(self.extraction.decisions)
            s["action_items"] = len(self.extraction.action_items)
        if self.validation:
            s["valid"] = self.validation.is_valid
            s["auto_fixes"] = self.validation.auto_fixes_applied
            s["errors"] = len(self.validation.errors)
            s["warnings"] = len(self.validation.warnings)
        if self.stored_ids:
            s["stored"] = {k: len(v) for k, v in self.stored_ids.items()}
        return s


def _create_llm(config: OntologyConfig) -> LLMProvider:
    """Create an LLM provider based on config."""
    provider = config.llm.provider.lower()
    if provider == "gemini":
        from ontology_engine.llm.gemini import GeminiProvider

        return GeminiProvider(config.llm)
    elif provider == "openai":
        from ontology_engine.llm.openai import OpenAIProvider

        return OpenAIProvider(config.llm)
    else:
        raise ValueError(f"Unknown LLM provider: {provider}")


def _create_llm_with_model(config: OntologyConfig, model: str) -> LLMProvider:
    """Create an LLM provider with an overridden model name."""
    from ontology_engine.core.config import LLMConfig

    new_llm_cfg = LLMConfig(
        provider=config.llm.provider,
        model=model,
        api_key=config.llm.api_key,
        temperature=config.llm.temperature,
        max_tokens=config.llm.max_tokens,
        timeout_seconds=config.llm.timeout_seconds,
        fast_model=config.llm.fast_model,
        default_model=model,
        strong_model=config.llm.strong_model,
    )
    new_config = OntologyConfig(
        llm=new_llm_cfg,
        database=config.database,
        pipeline=config.pipeline,
        known_entities=config.known_entities,
    )
    return _create_llm(new_config)
