"""Meta-learning engine: tracks intervention ROI and recommends best strategies."""
from __future__ import annotations

import enum
import json
import os
import sqlite3
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal
from uuid import uuid4

from logos_adk.persistence import (
    get_postgresql_schema_version,
    get_sqlite_schema_version,
    set_postgresql_schema_version,
    set_sqlite_schema_version,
)
from logos_adk.postgres_utils import import_psycopg, resolve_database_url

from .models import utcnow

META_LEARNING_SCHEMA_KIND = "meta_learning"
META_LEARNING_SCHEMA_VERSION = 1

InterventionSource = Literal["self_learning_loop", "manual", "knowledge_transfer", "rca"]


class InterventionType(str, enum.Enum):
    prompt_patch = "prompt_patch"
    routing_change = "routing_change"
    retrieval_change = "retrieval_change"
    tool_policy_change = "tool_policy_change"


class InterventionStatus(str, enum.Enum):
    applied = "applied"
    measured = "measured"
    rolled_back = "rolled_back"
    failed = "failed"


@dataclass
class InterventionRecord:
    intervention_id: str
    agent_name: str
    intervention_type: InterventionType
    source: InterventionSource
    trace_id: str | None = None
    context: dict[str, Any] = field(default_factory=dict)
    cost_estimate: float = 0.0
    created_at: str = field(default_factory=utcnow)
    uplift: float | None = None
    uplift_measured_at: str | None = None
    status: InterventionStatus = InterventionStatus.applied


@dataclass
class InterventionRecommendation:
    intervention_type: InterventionType | None
    action: str
    reason: str
    stats: dict[str, Any] = field(default_factory=dict)


@dataclass
class MetaLearningPolicy:
    min_interventions_for_stats: int = 5
    prefer_cheapest: bool = True
    auto_rollback_threshold: float = -0.05
    require_human_review_threshold: float = 0.02
    promote_threshold: float = 0.10
    roi_epsilon: float = 0.001


@dataclass
class ROIEntry:
    intervention_type: InterventionType
    avg_uplift: float
    avg_cost: float
    roi: float
    count: int
    success_rate: float


class MetaLearningEngine:
    """Tracks intervention outcomes and recommends best strategies by ROI."""

    def __init__(
        self,
        db_path: str | None = None,
        policy: MetaLearningPolicy | None = None,
        dsn: str | None = None,
        table_name: str = "meta_interventions",
    ) -> None:
        self.db_path = db_path or os.getenv(
            "LOGOS_ADK_META_LEARNING_DB_PATH",
            ".logos_adk/meta_learning.db",
        )
        self.dsn = resolve_database_url(
            dsn or os.getenv("LOGOS_ADK_META_LEARNING_DSN"),
            env_vars=("LOGOS_ADK_META_LEARNING_DSN", "LOGOS_ADK_DATABASE_URL", "DATABASE_URL"),
        )
        self.table_name = table_name
        self.policy = policy or MetaLearningPolicy()

    def _connect(self) -> sqlite3.Connection | Any:
        if self.dsn:
            return import_psycopg().connect(self.dsn)
        path = Path(self.db_path)
        if path.parent and not path.parent.exists():
            path.parent.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        self._ensure_schema(conn)
        return conn

    def _ensure_schema(self, conn: sqlite3.Connection | Any) -> None:
        if self.dsn:
            current = get_postgresql_schema_version(conn, META_LEARNING_SCHEMA_KIND) or 0
            with conn.cursor() as cursor:
                cursor.execute(
                    f"""
                    CREATE TABLE IF NOT EXISTS {self.table_name} (
                        intervention_id TEXT PRIMARY KEY,
                        agent_name TEXT NOT NULL,
                        intervention_type TEXT NOT NULL,
                        source TEXT NOT NULL,
                        trace_id TEXT,
                        context JSONB,
                        cost_estimate DOUBLE PRECISION DEFAULT 0.0,
                        created_at TEXT NOT NULL,
                        uplift DOUBLE PRECISION,
                        uplift_measured_at TEXT,
                        status TEXT NOT NULL DEFAULT 'applied'
                    )
                    """
                )
                cursor.execute(
                    f"""
                    CREATE INDEX IF NOT EXISTS idx_{self.table_name}_agent
                        ON {self.table_name} (agent_name, intervention_type, status)
                    """
                )
            if current < META_LEARNING_SCHEMA_VERSION:
                set_postgresql_schema_version(conn, META_LEARNING_SCHEMA_KIND, META_LEARNING_SCHEMA_VERSION)
            conn.commit()
            return

        current = get_sqlite_schema_version(conn, META_LEARNING_SCHEMA_KIND) or 0
        conn.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                intervention_id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                intervention_type TEXT NOT NULL,
                source TEXT NOT NULL,
                trace_id TEXT,
                context TEXT,
                cost_estimate REAL DEFAULT 0.0,
                created_at TEXT NOT NULL,
                uplift REAL,
                uplift_measured_at TEXT,
                status TEXT NOT NULL DEFAULT 'applied'
            )
            """
        )
        conn.execute(
            f"""
            CREATE INDEX IF NOT EXISTS idx_{self.table_name}_agent
                ON {self.table_name} (agent_name, intervention_type, status)
            """
        )
        if current < META_LEARNING_SCHEMA_VERSION:
            set_sqlite_schema_version(conn, META_LEARNING_SCHEMA_KIND, META_LEARNING_SCHEMA_VERSION)
        conn.commit()

    def _row_to_record(self, row: sqlite3.Row | Any) -> InterventionRecord:
        ctx = _row_value(row, "context", 5)
        return InterventionRecord(
            intervention_id=str(_row_value(row, "intervention_id", 0)),
            agent_name=str(_row_value(row, "agent_name", 1)),
            intervention_type=InterventionType(_row_value(row, "intervention_type", 2)),
            source=_row_value(row, "source", 3),
            trace_id=_row_value(row, "trace_id", 4),
            context=json.loads(ctx) if isinstance(ctx, str) and ctx else (ctx or {}),
            cost_estimate=float(_row_value(row, "cost_estimate", 6) or 0.0),
            created_at=str(_row_value(row, "created_at", 7)),
            uplift=_row_value(row, "uplift", 8),
            uplift_measured_at=_row_value(row, "uplift_measured_at", 9),
            status=InterventionStatus(_row_value(row, "status", 10)),
        )

    def log_intervention(
        self,
        agent_name: str,
        intervention_type: InterventionType,
        source: InterventionSource,
        trace_id: str | None = None,
        context: dict[str, Any] | None = None,
        cost_estimate: float = 0.0,
    ) -> InterventionRecord:
        record = InterventionRecord(
            intervention_id=str(uuid4()),
            agent_name=agent_name,
            intervention_type=intervention_type,
            source=source,
            trace_id=trace_id,
            context=context or {},
            cost_estimate=cost_estimate,
        )
        conn = self._connect()
        try:
            self._ensure_schema(conn)
            if self.dsn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f"""
                        INSERT INTO {self.table_name}
                            (intervention_id, agent_name, intervention_type, source,
                             trace_id, context, cost_estimate, created_at, status)
                        VALUES (%s, %s, %s, %s, %s, %s::jsonb, %s, %s, %s)
                        """,
                        (
                            record.intervention_id,
                            record.agent_name,
                            record.intervention_type.value,
                            record.source,
                            record.trace_id,
                            json.dumps(record.context),
                            record.cost_estimate,
                            record.created_at,
                            record.status.value,
                        ),
                    )
            else:
                conn.execute(
                    f"""
                    INSERT INTO {self.table_name}
                        (intervention_id, agent_name, intervention_type, source,
                         trace_id, context, cost_estimate, created_at, status)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        record.intervention_id,
                        record.agent_name,
                        record.intervention_type.value,
                        record.source,
                        record.trace_id,
                        json.dumps(record.context),
                        record.cost_estimate,
                        record.created_at,
                        record.status.value,
                    ),
                )
            conn.commit()
        finally:
            conn.close()
        return record

    def get_intervention(self, intervention_id: str) -> InterventionRecord | None:
        conn = self._connect()
        try:
            self._ensure_schema(conn)
            if self.dsn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f"SELECT * FROM {self.table_name} WHERE intervention_id = %s",
                        (intervention_id,),
                    )
                    row = cursor.fetchone()
            else:
                row = conn.execute(
                    f"SELECT * FROM {self.table_name} WHERE intervention_id = ?",
                    (intervention_id,),
                ).fetchone()
            return self._row_to_record(row) if row else None
        finally:
            conn.close()

    def measure_uplift(self, intervention_id: str, uplift: float) -> None:
        conn = self._connect()
        try:
            self._ensure_schema(conn)
            if self.dsn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f"""
                        UPDATE {self.table_name}
                        SET uplift = %s, uplift_measured_at = %s, status = 'measured'
                        WHERE intervention_id = %s
                        """,
                        (uplift, utcnow(), intervention_id),
                    )
            else:
                conn.execute(
                    f"""
                    UPDATE {self.table_name}
                    SET uplift = ?, uplift_measured_at = ?, status = 'measured'
                    WHERE intervention_id = ?
                    """,
                    (uplift, utcnow(), intervention_id),
                )
            conn.commit()
        finally:
            conn.close()

    def mark_rolled_back(self, intervention_id: str) -> None:
        conn = self._connect()
        try:
            self._ensure_schema(conn)
            if self.dsn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f"UPDATE {self.table_name} SET status = 'rolled_back' WHERE intervention_id = %s",
                        (intervention_id,),
                    )
            else:
                conn.execute(
                    f"UPDATE {self.table_name} SET status = 'rolled_back' WHERE intervention_id = ?",
                    (intervention_id,),
                )
            conn.commit()
        finally:
            conn.close()

    def intervention_stats(
        self,
        agent_name: str | None = None,
        intervention_type: InterventionType | None = None,
    ) -> dict[str, Any]:
        conn = self._connect()
        try:
            self._ensure_schema(conn)
            where_parts: list[str] = ["status = 'measured'"]
            params: list[Any] = []
            if agent_name:
                where_parts.append(f"agent_name = {'%s' if self.dsn else '?'}")
                params.append(agent_name)
            if intervention_type:
                where_parts.append(f"intervention_type = {'%s' if self.dsn else '?'}")
                params.append(intervention_type.value)
            where = " AND ".join(where_parts)
            query = f"""
                SELECT intervention_type,
                       COUNT(*) as cnt,
                       AVG(uplift) as avg_uplift,
                       AVG(cost_estimate) as avg_cost,
                       SUM(CASE WHEN uplift > 0 THEN 1 ELSE 0 END) as successes
                FROM {self.table_name}
                WHERE {where}
                GROUP BY intervention_type
                """
            if self.dsn:
                with conn.cursor() as cursor:
                    cursor.execute(query, params)
                    rows = cursor.fetchall()
            else:
                rows = conn.execute(query, params).fetchall()
            result: dict[str, Any] = {}
            for row in rows:
                cnt = _row_value(row, "cnt", 1)
                result[_row_value(row, "intervention_type", 0)] = {
                    "count": cnt,
                    "avg_uplift": _row_value(row, "avg_uplift", 2),
                    "avg_cost": _row_value(row, "avg_cost", 3),
                    "success_rate": (_row_value(row, "successes", 4) or 0) / cnt if cnt else 0.0,
                }
            return result
        finally:
            conn.close()

    def roi_leaderboard(self, agent_name: str | None = None) -> list[ROIEntry]:
        stats = self.intervention_stats(agent_name=agent_name)
        entries: list[ROIEntry] = []
        for itype_str, s in stats.items():
            avg_cost = s["avg_cost"]
            avg_uplift = s["avg_uplift"]
            roi = avg_uplift / max(avg_cost, self.policy.roi_epsilon)
            entries.append(
                ROIEntry(
                    intervention_type=InterventionType(itype_str),
                    avg_uplift=avg_uplift,
                    avg_cost=avg_cost,
                    roi=roi,
                    count=s["count"],
                    success_rate=s["success_rate"],
                )
            )
        entries.sort(key=lambda e: e.roi, reverse=True)
        return entries

    def recommend_intervention(
        self,
        agent_name: str,
        failure_types: list[str] | None = None,
        context: dict[str, Any] | None = None,
    ) -> InterventionRecommendation:
        stats = self.intervention_stats(agent_name=agent_name)

        # Filter by minimum sample size
        qualified: dict[str, Any] = {}
        for itype_str, s in stats.items():
            if s["count"] >= self.policy.min_interventions_for_stats:
                qualified[itype_str] = s

        if not qualified:
            return InterventionRecommendation(
                intervention_type=None,
                action="human_review",
                reason=f"Insufficient data: need at least {self.policy.min_interventions_for_stats} measured interventions per type",
                stats=stats,
            )

        # Rank by ROI, break ties by cost (lower is better when prefer_cheapest)
        def sort_key(item: tuple[str, dict]) -> tuple[float, float]:
            itype_str, s = item
            avg_cost = s["avg_cost"]
            avg_uplift = s["avg_uplift"]
            roi = avg_uplift / max(avg_cost, self.policy.roi_epsilon)
            cost_tiebreak = avg_cost if self.policy.prefer_cheapest else 0.0
            return (-roi, cost_tiebreak)

        ranked = sorted(qualified.items(), key=sort_key)
        best_type_str, best_stats = ranked[0]
        best_type = InterventionType(best_type_str)
        avg_uplift = best_stats["avg_uplift"]

        # Decision thresholds
        if avg_uplift < self.policy.auto_rollback_threshold:
            action = "rollback"
        elif avg_uplift < self.policy.require_human_review_threshold:
            action = "human_review"
        elif avg_uplift < self.policy.promote_threshold:
            action = "experiment"
        else:
            action = "promote"

        return InterventionRecommendation(
            intervention_type=best_type,
            action=action,
            reason=f"Best ROI type is {best_type.value} (avg_uplift={avg_uplift:.4f}, avg_cost={best_stats['avg_cost']:.4f}, count={best_stats['count']})",
            stats=best_stats,
        )

    def mark_failed(self, intervention_id: str) -> None:
        conn = self._connect()
        try:
            self._ensure_schema(conn)
            if self.dsn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f"UPDATE {self.table_name} SET status = 'failed' WHERE intervention_id = %s",
                        (intervention_id,),
                    )
            else:
                conn.execute(
                    f"UPDATE {self.table_name} SET status = 'failed' WHERE intervention_id = ?",
                    (intervention_id,),
                )
            conn.commit()
        finally:
            conn.close()


def _row_value(row: Any, name: str, index: int) -> Any:
    if row is None:
        return None
    try:
        return row[name]
    except (KeyError, TypeError, IndexError):
        return row[index]
