import json
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from sqlalchemy import text

from synthneura.core.logger import get_logger
from synthneura.core.schemas import ClinicalTrial
from synthneura.storage.base import ChangeRecord, RunMeta, StoreResult, TrialSink
from synthneura.storage.db import get_engine, resolve_db_url
from synthneura.storage.sqlite_sink import (
    _changed_fields,
    _core_to_dict,
    _raw_nct_id,
    _trial_core_values,
)

logger = get_logger(__name__)


def _jsonable_list(items: List[Any]) -> List[Any]:
    payload: List[Any] = []
    for item in items:
        if isinstance(item, dict):
            payload.append(item)
        elif hasattr(item, "model_dump"):
            payload.append(item.model_dump())
        elif hasattr(item, "dict"):
            payload.append(item.dict())
        else:
            payload.append(item)
    return payload


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def _ensure_schema(conn) -> None:
    conn.execute(
        text(
            """
            CREATE TABLE IF NOT EXISTS clinical_trials (
                nct_id TEXT PRIMARY KEY,
                title TEXT,
                status TEXT,
                phase TEXT,
                sponsor TEXT,
                disease_category TEXT,
                conditions TEXT,
                interventions TEXT,
                outcomes TEXT,
                raw_json TEXT,
                source TEXT,
                query TEXT,
                ingested_at TEXT,
                run_id TEXT
            )
            """
        )
    )
    conn.execute(
        text(
            "ALTER TABLE clinical_trials ADD COLUMN IF NOT EXISTS disease_category TEXT"
        )
    )

    conn.execute(
        text(
            """
            CREATE TABLE IF NOT EXISTS clinical_trials_history (
                id BIGSERIAL PRIMARY KEY,
                nct_id TEXT,
                title TEXT,
                status TEXT,
                phase TEXT,
                sponsor TEXT,
                disease_category TEXT,
                conditions TEXT,
                interventions TEXT,
                outcomes TEXT,
                raw_json TEXT,
                source TEXT,
                query TEXT,
                ingested_at TEXT,
                run_id TEXT,
                valid_from TEXT,
                valid_to TEXT
            )
            """
        )
    )
    conn.execute(
        text(
            "ALTER TABLE clinical_trials_history "
            "ADD COLUMN IF NOT EXISTS disease_category TEXT"
        )
    )

    conn.execute(
        text(
            """
            CREATE TABLE IF NOT EXISTS runs (
                run_id TEXT PRIMARY KEY,
                ingested_at TEXT,
                source TEXT,
                query TEXT,
                total INTEGER,
                new_count INTEGER,
                updated_count INTEGER,
                unchanged_count INTEGER,
                failed_count INTEGER
            )
            """
        )
    )

    conn.execute(
        text(
            """
            CREATE TABLE IF NOT EXISTS trial_diffs (
                id BIGSERIAL PRIMARY KEY,
                nct_id TEXT,
                run_id TEXT,
                ingested_at TEXT,
                changed_fields TEXT,
                before_json TEXT,
                after_json TEXT
            )
            """
        )
    )

    conn.execute(
        text(
            """
            CREATE TABLE IF NOT EXISTS trial_extractions (
                id BIGSERIAL PRIMARY KEY,
                nct_id TEXT,
                run_id TEXT,
                ingested_at TEXT,
                endpoints_json TEXT,
                biomarkers_json TEXT,
                provider TEXT,
                model TEXT,
                prompt_version TEXT,
                raw_response TEXT,
                validation_status TEXT,
                validation_errors TEXT,
                confidence DOUBLE PRECISION
            )
            """
        )
    )
    for column, col_type in (
        ("prompt_version", "TEXT"),
        ("raw_response", "TEXT"),
        ("validation_status", "TEXT"),
        ("validation_errors", "TEXT"),
        ("confidence", "DOUBLE PRECISION"),
    ):
        conn.execute(
            text(
                f"ALTER TABLE trial_extractions "
                f"ADD COLUMN IF NOT EXISTS {column} {col_type}"
            )
        )
    conn.execute(
        text(
            """
            CREATE UNIQUE INDEX IF NOT EXISTS idx_trial_extractions_nct_id_run_id
            ON trial_extractions(nct_id, run_id)
            """
        )
    )

    conn.execute(
        text(
            """
            CREATE TABLE IF NOT EXISTS trial_extraction_failures (
                id BIGSERIAL PRIMARY KEY,
                nct_id TEXT,
                run_id TEXT,
                ingested_at TEXT,
                provider TEXT,
                model TEXT,
                error TEXT
            )
            """
        )
    )
    conn.execute(
        text(
            """
            CREATE INDEX IF NOT EXISTS idx_trial_extraction_failures_run_id
            ON trial_extraction_failures(run_id)
            """
        )
    )

    conn.execute(
        text(
            """
            CREATE TABLE IF NOT EXISTS trial_extraction_diffs (
                id BIGSERIAL PRIMARY KEY,
                nct_id TEXT,
                run_id TEXT,
                ingested_at TEXT,
                endpoints_added TEXT,
                endpoints_removed TEXT,
                biomarkers_added TEXT,
                biomarkers_removed TEXT,
                prompt_version TEXT
            )
            """
        )
    )
    conn.execute(
        text(
            """
            CREATE INDEX IF NOT EXISTS idx_trial_extraction_diffs_run_id
            ON trial_extraction_diffs(run_id)
            """
        )
    )

    conn.execute(
        text(
            """
            CREATE INDEX IF NOT EXISTS idx_clinical_trials_history_nct_id_valid_from
            ON clinical_trials_history(nct_id, valid_from)
            """
        )
    )
    conn.execute(
        text(
            """
            CREATE INDEX IF NOT EXISTS idx_trial_diffs_nct_id_ingested_at
            ON trial_diffs(nct_id, ingested_at)
            """
        )
    )


class PostgresTrialSink(TrialSink):
    def __init__(self, db_url: Optional[str] = None) -> None:
        self.db_url = resolve_db_url(db_url)
        self.engine = get_engine(self.db_url)
        with self.engine.begin() as conn:
            _ensure_schema(conn)

    def fetch_previous_run_time(self, current_run_id: str) -> Optional[str]:
        with self.engine.begin() as conn:
            row = (
                conn.execute(
                    text(
                        """
                    SELECT ingested_at
                    FROM runs
                    WHERE run_id != :run_id
                    ORDER BY ingested_at DESC
                    LIMIT 1
                    """
                    ),
                    {"run_id": current_run_id},
                )
                .mappings()
                .first()
            )
            return row["ingested_at"] if row else None

    def fetch_diffs_between(
        self, start_time: Optional[str], end_time: Optional[str]
    ) -> List[ChangeRecord]:
        if not end_time:
            return []
        with self.engine.begin() as conn:
            if start_time:
                rows = (
                    conn.execute(
                        text(
                            """
                        SELECT nct_id, changed_fields, before_json, after_json
                        FROM trial_diffs
                        WHERE ingested_at > :start_time AND ingested_at <= :end_time
                        ORDER BY ingested_at DESC
                        """
                        ),
                        {"start_time": start_time, "end_time": end_time},
                    )
                    .mappings()
                    .all()
                )
            else:
                rows = (
                    conn.execute(
                        text(
                            """
                        SELECT nct_id, changed_fields, before_json, after_json
                        FROM trial_diffs
                        WHERE ingested_at <= :end_time
                        ORDER BY ingested_at DESC
                        """
                        ),
                        {"end_time": end_time},
                    )
                    .mappings()
                    .all()
                )

        changes: List[ChangeRecord] = []
        for row in rows:
            changes.append(
                ChangeRecord(
                    nct_id=row["nct_id"],
                    change_status="updated",
                    changed_fields=json.loads(row["changed_fields"] or "[]"),
                    before=json.loads(row["before_json"] or "{}"),
                    after=json.loads(row["after_json"] or "{}"),
                )
            )
        return changes

    def store_extractions(
        self,
        nct_id: str,
        endpoints: List[Any],
        biomarkers: List[Any],
        *,
        run_id: Optional[str] = None,
        ingested_at: Optional[str] = None,
        provider: Optional[str] = None,
        model: Optional[str] = None,
        prompt_version: Optional[str] = None,
        raw_response: Optional[str] = None,
        validation_status: Optional[str] = None,
        validation_errors: Optional[List[str]] = None,
        confidence: Optional[float] = None,
    ) -> None:
        with self.engine.begin() as conn:
            _ensure_schema(conn)
            endpoints_payload = _jsonable_list(endpoints)
            biomarkers_payload = _jsonable_list(biomarkers)
            conn.execute(
                text(
                    """
                    INSERT INTO trial_extractions
                    (
                        nct_id,
                        run_id,
                        ingested_at,
                        endpoints_json,
                        biomarkers_json,
                        provider,
                        model,
                        prompt_version,
                        raw_response,
                        validation_status,
                        validation_errors,
                        confidence
                    )
                    VALUES (:nct_id, :run_id, :ingested_at, :endpoints, :biomarkers,
                            :provider, :model, :prompt_version, :raw_response,
                            :validation_status, :validation_errors, :confidence)
                    ON CONFLICT (nct_id, run_id)
                    DO UPDATE SET
                        ingested_at = EXCLUDED.ingested_at,
                        endpoints_json = EXCLUDED.endpoints_json,
                        biomarkers_json = EXCLUDED.biomarkers_json,
                        provider = EXCLUDED.provider,
                        model = EXCLUDED.model,
                        prompt_version = EXCLUDED.prompt_version,
                        raw_response = EXCLUDED.raw_response,
                        validation_status = EXCLUDED.validation_status,
                        validation_errors = EXCLUDED.validation_errors,
                        confidence = EXCLUDED.confidence
                    """
                ),
                {
                    "nct_id": nct_id,
                    "run_id": run_id,
                    "ingested_at": ingested_at,
                    "endpoints": json.dumps(endpoints_payload, ensure_ascii=True),
                    "biomarkers": json.dumps(biomarkers_payload, ensure_ascii=True),
                    "provider": provider,
                    "model": model,
                    "prompt_version": prompt_version,
                    "raw_response": raw_response,
                    "validation_status": validation_status,
                    "validation_errors": json.dumps(
                        validation_errors or [], ensure_ascii=True
                    ),
                    "confidence": confidence,
                },
            )

    def store_extraction_failure(
        self,
        nct_id: str,
        *,
        run_id: Optional[str] = None,
        ingested_at: Optional[str] = None,
        provider: Optional[str] = None,
        model: Optional[str] = None,
        error: Optional[str] = None,
    ) -> None:
        with self.engine.begin() as conn:
            _ensure_schema(conn)
            conn.execute(
                text(
                    """
                    INSERT INTO trial_extraction_failures
                    (
                        nct_id,
                        run_id,
                        ingested_at,
                        provider,
                        model,
                        error
                    )
                    VALUES (:nct_id, :run_id, :ingested_at, :provider, :model, :error)
                    """
                ),
                {
                    "nct_id": nct_id,
                    "run_id": run_id,
                    "ingested_at": ingested_at,
                    "provider": provider,
                    "model": model,
                    "error": error,
                },
            )

    def fetch_previous_extraction(
        self, nct_id: str, current_run_id: str
    ) -> Optional[Dict[str, Any]]:
        with self.engine.begin() as conn:
            row = (
                conn.execute(
                    text(
                        """
                    SELECT endpoints_json, biomarkers_json, run_id, ingested_at
                    FROM trial_extractions
                    WHERE nct_id = :nct_id AND run_id != :run_id
                    ORDER BY ingested_at DESC
                    LIMIT 1
                    """
                    ),
                    {"nct_id": nct_id, "run_id": current_run_id},
                )
                .mappings()
                .first()
            )
            return dict(row) if row else None

    def store_extraction_diff(
        self,
        nct_id: str,
        *,
        run_id: str,
        ingested_at: str,
        endpoints_added: List[str],
        endpoints_removed: List[str],
        biomarkers_added: List[str],
        biomarkers_removed: List[str],
        prompt_version: Optional[str] = None,
    ) -> None:
        with self.engine.begin() as conn:
            _ensure_schema(conn)
            conn.execute(
                text(
                    """
                    INSERT INTO trial_extraction_diffs
                    (
                        nct_id,
                        run_id,
                        ingested_at,
                        endpoints_added,
                        endpoints_removed,
                        biomarkers_added,
                        biomarkers_removed,
                        prompt_version
                    )
                    VALUES (:nct_id, :run_id, :ingested_at, :endpoints_added,
                            :endpoints_removed, :biomarkers_added,
                            :biomarkers_removed, :prompt_version)
                    """
                ),
                {
                    "nct_id": nct_id,
                    "run_id": run_id,
                    "ingested_at": ingested_at,
                    "endpoints_added": json.dumps(endpoints_added, ensure_ascii=True),
                    "endpoints_removed": json.dumps(
                        endpoints_removed, ensure_ascii=True
                    ),
                    "biomarkers_added": json.dumps(biomarkers_added, ensure_ascii=True),
                    "biomarkers_removed": json.dumps(
                        biomarkers_removed, ensure_ascii=True
                    ),
                    "prompt_version": prompt_version,
                },
            )

    def store_trials(
        self,
        trials: List[ClinicalTrial],
        raw_trials: List[Dict[str, Any]],
        run_meta: RunMeta,
    ) -> StoreResult:
        logger.debug("Storing %d trials in Postgres db=%s", len(trials), self.db_url)

        raw_by_id = {}
        for raw_trial in raw_trials:
            nct_id = _raw_nct_id(raw_trial)
            if nct_id:
                raw_by_id[nct_id] = raw_trial

        counts = {"new": 0, "updated": 0, "unchanged": 0, "failed": 0}
        changes: List[ChangeRecord] = []
        with self.engine.begin() as conn:
            _ensure_schema(conn)
            for trial in trials:
                existing = (
                    conn.execute(
                        text(
                            """
                        SELECT
                            title,
                            status,
                            phase,
                            sponsor,
                            disease_category,
                            conditions,
                            interventions,
                            outcomes,
                            raw_json,
                            source,
                            query,
                            ingested_at,
                            run_id
                        FROM clinical_trials
                        WHERE nct_id = :nct_id
                        """
                        ),
                        {"nct_id": trial.nct_id},
                    )
                    .mappings()
                    .first()
                )
                existing_core = (
                    (
                        existing["title"],
                        existing["status"],
                        existing["phase"],
                        existing["sponsor"],
                        existing["disease_category"],
                        existing["conditions"],
                        existing["interventions"],
                        existing["outcomes"],
                    )
                    if existing
                    else None
                )
                existing_raw_json = existing["raw_json"] if existing else None
                existing_source = existing["source"] if existing else None
                existing_query = existing["query"] if existing else None
                existing_ingested_at = existing["ingested_at"] if existing else None
                existing_run_id = existing["run_id"] if existing else None
                new_core = _trial_core_values(trial)

                if existing is None:
                    change_status = "new"
                elif existing_core == new_core:
                    change_status = "unchanged"
                else:
                    change_status = "updated"

                ingested_at_value = run_meta.get("ingested_at") or _utc_now()
                if change_status == "unchanged" and existing_ingested_at:
                    ingested_at_value = existing_ingested_at
                raw_json = json.dumps(
                    raw_by_id.get(trial.nct_id, {}), ensure_ascii=True
                )
                if trial.nct_id not in raw_by_id and existing_raw_json is not None:
                    raw_json = existing_raw_json
                if change_status == "updated" and existing is not None:
                    changed_fields = _changed_fields(
                        existing_core or new_core, new_core
                    )
                    before = _core_to_dict(existing_core or new_core)
                    after = _core_to_dict(new_core)
                    changes.append(
                        ChangeRecord(
                            nct_id=trial.nct_id,
                            change_status=change_status,
                            changed_fields=changed_fields,
                            before=before,
                            after=after,
                        )
                    )
                    conn.execute(
                        text(
                            """
                            INSERT INTO trial_diffs
                            (
                                nct_id,
                                run_id,
                                ingested_at,
                                changed_fields,
                                before_json,
                                after_json
                            )
                            VALUES (:nct_id, :run_id, :ingested_at, :changed_fields,
                                    :before_json, :after_json)
                            """
                        ),
                        {
                            "nct_id": trial.nct_id,
                            "run_id": run_meta.get("run_id"),
                            "ingested_at": ingested_at_value,
                            "changed_fields": json.dumps(
                                changed_fields, ensure_ascii=True
                            ),
                            "before_json": json.dumps(before, ensure_ascii=True),
                            "after_json": json.dumps(after, ensure_ascii=True),
                        },
                    )
                    (
                        existing_title,
                        existing_status,
                        existing_phase,
                        existing_sponsor,
                        existing_disease_category,
                        existing_conditions,
                        existing_interventions,
                        existing_outcomes,
                    ) = (
                        existing_core or new_core
                    )
                    conn.execute(
                        text(
                            """
                            INSERT INTO clinical_trials_history
                            (
                                nct_id,
                                title,
                                status,
                                phase,
                                sponsor,
                                disease_category,
                                conditions,
                                interventions,
                                outcomes,
                                raw_json,
                                source,
                                query,
                                ingested_at,
                                run_id,
                                valid_from,
                                valid_to
                            )
                            VALUES (
                                :nct_id,
                                :title,
                                :status,
                                :phase,
                                :sponsor,
                                :disease_category,
                                :conditions,
                                :interventions,
                                :outcomes,
                                :raw_json,
                                :source,
                                :query,
                                :ingested_at,
                                :run_id,
                                :valid_from,
                                :valid_to
                            )
                            """
                        ),
                        {
                            "nct_id": trial.nct_id,
                            "title": existing_title,
                            "status": existing_status,
                            "phase": existing_phase,
                            "sponsor": existing_sponsor,
                            "disease_category": existing_disease_category,
                            "conditions": existing_conditions,
                            "interventions": existing_interventions,
                            "outcomes": existing_outcomes,
                            "raw_json": existing_raw_json,
                            "source": existing_source,
                            "query": existing_query,
                            "ingested_at": existing_ingested_at,
                            "run_id": existing_run_id,
                            "valid_from": existing_ingested_at,
                            "valid_to": ingested_at_value,
                        },
                    )
                if change_status != "unchanged":
                    conn.execute(
                        text(
                            """
                            INSERT INTO clinical_trials
                            (
                                nct_id,
                                title,
                                status,
                                phase,
                                sponsor,
                                disease_category,
                                conditions,
                                interventions,
                                outcomes,
                                raw_json,
                                source,
                                query,
                                ingested_at,
                                run_id
                            )
                            VALUES (
                                :nct_id,
                                :title,
                                :status,
                                :phase,
                                :sponsor,
                                :disease_category,
                                :conditions,
                                :interventions,
                                :outcomes,
                                :raw_json,
                                :source,
                                :query,
                                :ingested_at,
                                :run_id
                            )
                            ON CONFLICT (nct_id)
                            DO UPDATE SET
                                title = EXCLUDED.title,
                                status = EXCLUDED.status,
                                phase = EXCLUDED.phase,
                                sponsor = EXCLUDED.sponsor,
                                disease_category = EXCLUDED.disease_category,
                                conditions = EXCLUDED.conditions,
                                interventions = EXCLUDED.interventions,
                                outcomes = EXCLUDED.outcomes,
                                raw_json = EXCLUDED.raw_json,
                                source = EXCLUDED.source,
                                query = EXCLUDED.query,
                                ingested_at = EXCLUDED.ingested_at,
                                run_id = EXCLUDED.run_id
                            """
                        ),
                        {
                            "nct_id": trial.nct_id,
                            "title": trial.title,
                            "status": trial.status,
                            "phase": trial.phase,
                            "sponsor": trial.sponsor,
                            "disease_category": trial.disease_category,
                            "conditions": ",".join(trial.conditions or []),
                            "interventions": ",".join(trial.interventions or []),
                            "outcomes": ",".join(trial.outcomes or []),
                            "raw_json": raw_json,
                            "source": run_meta.get("source"),
                            "query": run_meta.get("query"),
                            "ingested_at": ingested_at_value,
                            "run_id": run_meta.get("run_id"),
                        },
                    )
                else:
                    if run_meta.get("run_id"):
                        conn.execute(
                            text(
                                """
                                UPDATE clinical_trials
                                SET run_id = :run_id, source = :source, query = :query
                                WHERE nct_id = :nct_id
                                """
                            ),
                            {
                                "run_id": run_meta.get("run_id"),
                                "source": run_meta.get("source"),
                                "query": run_meta.get("query"),
                                "nct_id": trial.nct_id,
                            },
                        )
                counts[change_status] += 1

            run_id = run_meta.get("run_id")
            if run_id:
                conn.execute(
                    text(
                        """
                        INSERT INTO runs
                        (
                            run_id,
                            ingested_at,
                            source,
                            query,
                            total,
                            new_count,
                            updated_count,
                            unchanged_count,
                            failed_count
                        )
                        VALUES (
                            :run_id,
                            :ingested_at,
                            :source,
                            :query,
                            :total,
                            :new_count,
                            :updated_count,
                            :unchanged_count,
                            :failed_count
                        )
                        ON CONFLICT (run_id)
                        DO UPDATE SET
                            ingested_at = EXCLUDED.ingested_at,
                            source = EXCLUDED.source,
                            query = EXCLUDED.query,
                            total = EXCLUDED.total,
                            new_count = EXCLUDED.new_count,
                            updated_count = EXCLUDED.updated_count,
                            unchanged_count = EXCLUDED.unchanged_count,
                            failed_count = EXCLUDED.failed_count
                        """
                    ),
                    {
                        "run_id": run_id,
                        "ingested_at": run_meta.get("ingested_at"),
                        "source": run_meta.get("source"),
                        "query": run_meta.get("query"),
                        "total": len(trials),
                        "new_count": counts["new"],
                        "updated_count": counts["updated"],
                        "unchanged_count": counts["unchanged"],
                        "failed_count": counts["failed"],
                    },
                )

        return StoreResult(counts=counts, changes=changes)
