"""Model routing learning for adaptive provider selection."""
from __future__ import annotations

from dataclasses import dataclass, field
from fnmatch import fnmatch
import hashlib
import os
from pathlib import Path
import sqlite3
from typing import Any
from uuid import uuid4

from logos_adk.cost import DEFAULT_RATES
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 FeedbackRecord, OutcomeRecord, utcnow
from .store import LearningStore


MODEL_ROUTING_SCHEMA_KIND = "model_routing_learning"
MODEL_ROUTING_SCHEMA_VERSION = 1


ConfidenceTier = str


def _segment(value: str | None) -> str:
    return value or ""


def _quality_from_feedback(feedback: list[FeedbackRecord], outcomes: list[OutcomeRecord]) -> float:
    score = 0.0
    count = 0
    for item in feedback:
        if item.kind == "thumbs_up":
            score += 1.0
            count += 1
        elif item.kind == "thumbs_down":
            score -= 1.0
            count += 1
        elif item.kind == "numeric_rating" and item.rating is not None:
            score += max(-1.0, min(1.0, (float(item.rating) - 3.0) / 2.0))
            count += 1
        elif item.kind == "edited_answer":
            score -= 0.75
            count += 1
        elif item.failure_types:
            score -= 0.5
            count += 1
    for item in outcomes:
        if item.kind in {"conversion", "resolution"} and bool(item.value):
            score += 1.0
            count += 1
        elif item.kind == "reply_rate":
            score += float(item.value or 0.0)
            count += 1
        elif item.kind == "ctr":
            score += min(1.0, float(item.value or 0.0))
            count += 1
        elif item.kind in {"escalation", "human_takeover"} and bool(item.value):
            score -= 1.0
            count += 1
    if count <= 0:
        return 0.0
    return max(-1.0, min(1.0, score / count))


def _bayes_rate(successes: int, failures: int) -> float:
    total = successes + failures
    if total <= 0:
        return 0.5
    return (successes + 1.0) / (total + 2.0)


@dataclass
class RoutingBudget:
    """Budget policy for a routing segment."""

    workspace_id: str | None = None
    task_type: str | None = None
    soft_limit_usd: float = 0.0
    hard_limit_usd: float | None = None
    threshold_ratio: float = 0.8
    period: str = "daily"
    updated_at: str = field(default_factory=utcnow)


@dataclass
class CanaryDeployment:
    """Partial rollout of a candidate model for live comparison."""

    id: str
    model_key: str
    rollout_percentage: int
    workspace_id: str | None = None
    task_type: str | None = None
    user_segment: str | None = None
    status: str = "active"
    created_at: str = field(default_factory=utcnow)
    ended_at: str | None = None


@dataclass
class ModelScore:
    """Aggregated routing leaderboard entry."""

    model_key: str
    workspace_id: str | None = None
    task_type: str | None = None
    confidence_tier: str | None = None
    success_count: int = 0
    failure_count: int = 0
    average_latency_ms: float = 0.0
    average_cost_usd: float = 0.0
    average_quality_score: float = 0.0
    win_rate: float = 0.0
    score: float = 0.0
    canary_calls: int = 0
    last_updated: str = field(default_factory=utcnow)


@dataclass
class RoutingDecision:
    """Chosen ordering for a request."""

    ordered_model_keys: list[str]
    selected_model_key: str
    confidence_tier: str
    budget_ratio: float = 0.0
    canary_model_key: str | None = None
    used_canary: bool = False
    routed_for_budget: bool = False


class ModelRoutingLearningEngine:
    """Persistent model routing engine with budgets and canaries."""

    def __init__(
        self,
        path: str | None = None,
        table_prefix: str | None = None,
        dsn: str | None = None,
    ) -> None:
        self.path = path or os.getenv(
            "LOGOS_ADK_MODEL_ROUTING_DB_PATH",
            ".logos_adk/model_routing_learning.db",
        )
        self.dsn = resolve_database_url(
            dsn or os.getenv("LOGOS_ADK_MODEL_ROUTING_DSN"),
            env_vars=("LOGOS_ADK_MODEL_ROUTING_DSN", "LOGOS_ADK_DATABASE_URL", "DATABASE_URL"),
        )
        self.table_prefix = table_prefix or os.getenv(
            "LOGOS_ADK_MODEL_ROUTING_TABLE_PREFIX",
            "model_routing_learning",
        )

    def _table(self, suffix: str) -> str:
        return f"{self.table_prefix}_{suffix}"

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

    def _ensure_schema(self, connection: sqlite3.Connection | Any) -> None:
        if self.dsn:
            current = get_postgresql_schema_version(connection, MODEL_ROUTING_SCHEMA_KIND) or 0
            connection.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self._table('events')} (
                    id TEXT PRIMARY KEY,
                    trace_id TEXT,
                    workspace_id TEXT NOT NULL,
                    task_type TEXT NOT NULL,
                    user_segment TEXT NOT NULL,
                    model_key TEXT NOT NULL,
                    confidence_tier TEXT NOT NULL,
                    success INTEGER NOT NULL,
                    latency_ms DOUBLE PRECISION NOT NULL,
                    cost_usd DOUBLE PRECISION NOT NULL,
                    quality_score DOUBLE PRECISION NOT NULL DEFAULT 0,
                    was_selected INTEGER NOT NULL DEFAULT 1,
                    was_canary INTEGER NOT NULL DEFAULT 0,
                    synced_at TEXT,
                    created_at TEXT NOT NULL
                )
                """
            )
            connection.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self._table('events')}_segment ON {self._table('events')} (workspace_id, task_type, confidence_tier, model_key, created_at DESC)"
            )
            connection.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self._table('events')}_trace ON {self._table('events')} (trace_id, created_at DESC)"
            )
            connection.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self._table('budgets')} (
                    workspace_id TEXT NOT NULL,
                    task_type TEXT NOT NULL,
                    soft_limit_usd DOUBLE PRECISION NOT NULL,
                    hard_limit_usd DOUBLE PRECISION,
                    threshold_ratio DOUBLE PRECISION NOT NULL,
                    period TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    PRIMARY KEY (workspace_id, task_type)
                )
                """
            )
            connection.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self._table('canaries')} (
                    id TEXT PRIMARY KEY,
                    model_key TEXT NOT NULL,
                    rollout_percentage INTEGER NOT NULL,
                    workspace_id TEXT NOT NULL,
                    task_type TEXT NOT NULL,
                    user_segment TEXT NOT NULL,
                    status TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    ended_at TEXT
                )
                """
            )
            connection.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self._table('canaries')}_lookup ON {self._table('canaries')} (status, workspace_id, task_type, user_segment, created_at DESC)"
            )
            if current < MODEL_ROUTING_SCHEMA_VERSION:
                set_postgresql_schema_version(
                    connection,
                    MODEL_ROUTING_SCHEMA_KIND,
                    MODEL_ROUTING_SCHEMA_VERSION,
                )
            connection.commit()
            return
        current = get_sqlite_schema_version(connection, MODEL_ROUTING_SCHEMA_KIND) or 0
        connection.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self._table('events')} (
                id TEXT PRIMARY KEY,
                trace_id TEXT,
                workspace_id TEXT NOT NULL,
                task_type TEXT NOT NULL,
                user_segment TEXT NOT NULL,
                model_key TEXT NOT NULL,
                confidence_tier TEXT NOT NULL,
                success INTEGER NOT NULL,
                latency_ms REAL NOT NULL,
                cost_usd REAL NOT NULL,
                quality_score REAL NOT NULL DEFAULT 0,
                was_selected INTEGER NOT NULL DEFAULT 1,
                was_canary INTEGER NOT NULL DEFAULT 0,
                synced_at TEXT,
                created_at TEXT NOT NULL
            )
            """
        )
        connection.execute(
            f"CREATE INDEX IF NOT EXISTS idx_{self._table('events')}_segment ON {self._table('events')} (workspace_id, task_type, confidence_tier, model_key, created_at DESC)"
        )
        connection.execute(
            f"CREATE INDEX IF NOT EXISTS idx_{self._table('events')}_trace ON {self._table('events')} (trace_id, created_at DESC)"
        )
        connection.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self._table('budgets')} (
                workspace_id TEXT NOT NULL,
                task_type TEXT NOT NULL,
                soft_limit_usd REAL NOT NULL,
                hard_limit_usd REAL,
                threshold_ratio REAL NOT NULL,
                period TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                PRIMARY KEY (workspace_id, task_type)
            )
            """
        )
        connection.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self._table('canaries')} (
                id TEXT PRIMARY KEY,
                model_key TEXT NOT NULL,
                rollout_percentage INTEGER NOT NULL,
                workspace_id TEXT NOT NULL,
                task_type TEXT NOT NULL,
                user_segment TEXT NOT NULL,
                status TEXT NOT NULL,
                created_at TEXT NOT NULL,
                ended_at TEXT
            )
            """
        )
        connection.execute(
            f"CREATE INDEX IF NOT EXISTS idx_{self._table('canaries')}_lookup ON {self._table('canaries')} (status, workspace_id, task_type, user_segment, created_at DESC)"
        )
        if current < MODEL_ROUTING_SCHEMA_VERSION:
            set_sqlite_schema_version(
                connection,
                MODEL_ROUTING_SCHEMA_KIND,
                MODEL_ROUTING_SCHEMA_VERSION,
            )
        connection.commit()

    def record_route(
        self,
        *,
        model_key: str,
        success: bool,
        latency_ms: float,
        cost_usd: float,
        confidence_tier: str,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
        trace_id: str | None = None,
        was_canary: bool = False,
        was_selected: bool = True,
        quality_score: float = 0.0,
    ) -> None:
        with self._connect() as connection:
            connection.execute(
                f"""
                INSERT INTO {self._table('events')}
                (id, trace_id, workspace_id, task_type, user_segment, model_key,
                 confidence_tier, success, latency_ms, cost_usd, quality_score,
                 was_selected, was_canary, synced_at, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    uuid4().hex,
                    trace_id,
                    _segment(workspace_id),
                    _segment(task_type),
                    _segment(user_segment),
                    model_key,
                    confidence_tier,
                    1 if success else 0,
                    max(0.0, float(latency_ms)),
                    max(0.0, float(cost_usd)),
                    max(-1.0, min(1.0, float(quality_score))),
                    1 if was_selected else 0,
                    1 if was_canary else 0,
                    None,
                    utcnow(),
                ),
            )

    def set_budget(
        self,
        *,
        soft_limit_usd: float,
        workspace_id: str | None = None,
        task_type: str | None = None,
        hard_limit_usd: float | None = None,
        threshold_ratio: float = 0.8,
        period: str = "daily",
    ) -> RoutingBudget:
        record = RoutingBudget(
            workspace_id=workspace_id,
            task_type=task_type,
            soft_limit_usd=max(0.0, float(soft_limit_usd)),
            hard_limit_usd=None if hard_limit_usd is None else max(0.0, float(hard_limit_usd)),
            threshold_ratio=max(0.0, min(1.0, float(threshold_ratio))),
            period=period,
        )
        with self._connect() as connection:
            connection.execute(
                f"""
                INSERT INTO {self._table('budgets')}
                (workspace_id, task_type, soft_limit_usd, hard_limit_usd, threshold_ratio, period, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(workspace_id, task_type) DO UPDATE SET
                    soft_limit_usd = excluded.soft_limit_usd,
                    hard_limit_usd = excluded.hard_limit_usd,
                    threshold_ratio = excluded.threshold_ratio,
                    period = excluded.period,
                    updated_at = excluded.updated_at
                """,
                (
                    _segment(workspace_id),
                    _segment(task_type),
                    record.soft_limit_usd,
                    record.hard_limit_usd,
                    record.threshold_ratio,
                    record.period,
                    record.updated_at,
                ),
            )
        return record

    def get_budget(
        self,
        *,
        workspace_id: str | None = None,
        task_type: str | None = None,
    ) -> RoutingBudget | None:
        with self._connect() as connection:
            row = connection.execute(
                f"""
                SELECT * FROM {self._table('budgets')}
                WHERE workspace_id = ? AND task_type = ?
                """,
                (_segment(workspace_id), _segment(task_type)),
            ).fetchone()
        if row is None:
            return None
        return RoutingBudget(
            workspace_id=workspace_id,
            task_type=task_type,
            soft_limit_usd=float(row["soft_limit_usd"]),
            hard_limit_usd=float(row["hard_limit_usd"]) if row["hard_limit_usd"] is not None else None,
            threshold_ratio=float(row["threshold_ratio"]),
            period=str(row["period"]),
            updated_at=str(row["updated_at"]),
        )

    def budget_status(
        self,
        *,
        workspace_id: str | None = None,
        task_type: str | None = None,
    ) -> dict[str, Any]:
        budget = self.get_budget(workspace_id=workspace_id, task_type=task_type)
        spent = self._spent_cost(workspace_id=workspace_id, task_type=task_type)
        soft_limit = float(budget.soft_limit_usd) if budget is not None else 0.0
        hard_limit = float(budget.hard_limit_usd) if budget is not None and budget.hard_limit_usd is not None else None
        ratio = spent / soft_limit if soft_limit > 0 else 0.0
        return {
            "workspace_id": workspace_id,
            "task_type": task_type,
            "spent_usd": round(spent, 8),
            "soft_limit_usd": soft_limit,
            "hard_limit_usd": hard_limit,
            "ratio": round(ratio, 6),
            "near_limit": budget is not None and ratio >= budget.threshold_ratio,
            "hard_limit_reached": hard_limit is not None and spent >= hard_limit,
        }

    def register_canary(
        self,
        *,
        model_key: str,
        rollout_percentage: int,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
    ) -> CanaryDeployment:
        record = CanaryDeployment(
            id=uuid4().hex,
            model_key=model_key,
            rollout_percentage=max(1, min(100, int(rollout_percentage))),
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
        )
        with self._connect() as connection:
            connection.execute(
                f"""
                INSERT INTO {self._table('canaries')}
                (id, model_key, rollout_percentage, workspace_id, task_type, user_segment, status, created_at, ended_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    record.id,
                    record.model_key,
                    record.rollout_percentage,
                    _segment(workspace_id),
                    _segment(task_type),
                    _segment(user_segment),
                    record.status,
                    record.created_at,
                    record.ended_at,
                ),
            )
        return record

    def list_canaries(
        self,
        *,
        status: str | None = "active",
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
    ) -> list[CanaryDeployment]:
        clauses: list[str] = []
        params: list[Any] = []
        if status is not None:
            clauses.append("status = ?")
            params.append(status)
        if workspace_id is not None:
            clauses.append("workspace_id = ?")
            params.append(_segment(workspace_id))
        if task_type is not None:
            clauses.append("task_type = ?")
            params.append(_segment(task_type))
        if user_segment is not None:
            clauses.append("user_segment = ?")
            params.append(_segment(user_segment))
        where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        with self._connect() as connection:
            rows = connection.execute(
                f"SELECT * FROM {self._table('canaries')} {where_sql} ORDER BY created_at DESC",
                tuple(params),
            ).fetchall()
        return [
            CanaryDeployment(
                id=str(row["id"]),
                model_key=str(row["model_key"]),
                rollout_percentage=int(row["rollout_percentage"]),
                workspace_id=None if str(row["workspace_id"]) == "" else str(row["workspace_id"]),
                task_type=None if str(row["task_type"]) == "" else str(row["task_type"]),
                user_segment=None if str(row["user_segment"]) == "" else str(row["user_segment"]),
                status=str(row["status"]),
                created_at=str(row["created_at"]),
                ended_at=str(row["ended_at"]) if row["ended_at"] is not None else None,
            )
            for row in rows
        ]

    def list_model_scores(
        self,
        *,
        workspace_id: str | None = None,
        task_type: str | None = None,
        confidence_tier: str | None = None,
    ) -> list[ModelScore]:
        clauses = ["workspace_id = ?"]
        params: list[Any] = [_segment(workspace_id)]
        if task_type is not None:
            clauses.append("task_type = ?")
            params.append(_segment(task_type))
        if confidence_tier is not None:
            clauses.append("confidence_tier = ?")
            params.append(confidence_tier)
        query = f"""
            SELECT
                model_key,
                confidence_tier,
                SUM(CASE WHEN success = 1 THEN 1 ELSE 0 END) AS success_count,
                SUM(CASE WHEN success = 0 THEN 1 ELSE 0 END) AS failure_count,
                AVG(latency_ms) AS average_latency_ms,
                AVG(cost_usd) AS average_cost_usd,
                AVG(quality_score) AS average_quality_score,
                SUM(CASE WHEN was_canary = 1 THEN 1 ELSE 0 END) AS canary_calls,
                MAX(created_at) AS last_updated
            FROM {self._table('events')}
            WHERE {' AND '.join(clauses)}
            GROUP BY workspace_id, task_type, model_key, confidence_tier
        """
        with self._connect() as connection:
            rows = connection.execute(query, tuple(params)).fetchall()
        items: list[ModelScore] = []
        for row in rows:
            success_count = int(row["success_count"] or 0)
            failure_count = int(row["failure_count"] or 0)
            avg_latency = float(row["average_latency_ms"] or 0.0)
            avg_cost = float(row["average_cost_usd"] or 0.0)
            avg_quality = float(row["average_quality_score"] or 0.0)
            win_rate = success_count / max(1, success_count + failure_count)
            latency_component = 1.0 / (1.0 + (avg_latency / 1000.0))
            cost_component = 1.0 / (1.0 + (avg_cost * 200.0))
            score = (
                ((avg_quality + 1.0) / 2.0) * 0.45
                + _bayes_rate(success_count, failure_count) * 0.25
                + latency_component * 0.15
                + cost_component * 0.15
            )
            items.append(
                ModelScore(
                    model_key=str(row["model_key"]),
                    workspace_id=workspace_id,
                    task_type=task_type,
                    confidence_tier=str(row["confidence_tier"]),
                    success_count=success_count,
                    failure_count=failure_count,
                    average_latency_ms=avg_latency,
                    average_cost_usd=avg_cost,
                    average_quality_score=avg_quality,
                    win_rate=win_rate,
                    score=score,
                    canary_calls=int(row["canary_calls"] or 0),
                    last_updated=str(row["last_updated"] or utcnow()),
                )
            )
        items.sort(key=lambda item: (-item.score, -item.average_quality_score, item.average_cost_usd, item.model_key))
        return items

    def route_candidates(
        self,
        model_keys: list[str],
        *,
        prompt: str | None = None,
        assignment_key: str | None = None,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
    ) -> RoutingDecision:
        keys = list(dict.fromkeys(model_keys))
        if not keys:
            raise ValueError("route_candidates requires at least one model key")
        confidence_tier = self.estimate_confidence(prompt or "", task_type=task_type)
        scores = {
            item.model_key: item
            for item in self.list_model_scores(
                workspace_id=workspace_id,
                task_type=task_type,
                confidence_tier=confidence_tier,
            )
        }
        budget = self.budget_status(workspace_id=workspace_id, task_type=task_type)
        cheapest_key = min(keys, key=self._estimated_cost_for_model)

        if confidence_tier == "high":
            ordered = sorted(
                keys,
                key=lambda key: (
                    -(scores.get(key).average_quality_score if key in scores else 0.0),
                    -(scores.get(key).score if key in scores else 0.5),
                    self._estimated_cost_for_model(key),
                ),
            )
        elif confidence_tier == "low":
            ordered = sorted(
                keys,
                key=lambda key: (
                    self._estimated_cost_for_model(key),
                    -(scores.get(key).score if key in scores else 0.5),
                ),
            )
        else:
            ordered = sorted(
                keys,
                key=lambda key: (
                    -(scores.get(key).score if key in scores else 0.5),
                    self._estimated_cost_for_model(key),
                ),
            )

        routed_for_budget = False
        if budget["hard_limit_reached"] or budget["near_limit"]:
            ordered = [cheapest_key] + [key for key in ordered if key != cheapest_key]
            routed_for_budget = True

        canary_model_key: str | None = None
        used_canary = False
        canary = self._active_canary(
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            candidate_keys=ordered,
        )
        if canary is not None:
            canary_model_key = canary.model_key
            if self._should_assign_canary(canary, assignment_key):
                ordered = [canary.model_key] + [key for key in ordered if key != canary.model_key]
                used_canary = True

        return RoutingDecision(
            ordered_model_keys=ordered,
            selected_model_key=ordered[0],
            confidence_tier=confidence_tier,
            budget_ratio=float(budget["ratio"]),
            canary_model_key=canary_model_key,
            used_canary=used_canary,
            routed_for_budget=routed_for_budget,
        )

    def inspect_routing_decision(
        self,
        model_keys: list[str],
        *,
        prompt: str | None = None,
        assignment_key: str | None = None,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
    ) -> dict[str, Any]:
        decision = self.route_candidates(
            model_keys,
            prompt=prompt,
            assignment_key=assignment_key,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
        )
        scores = self.list_model_scores(
            workspace_id=workspace_id,
            task_type=task_type,
            confidence_tier=decision.confidence_tier,
        )
        return {
            "decision": {
                "ordered_model_keys": decision.ordered_model_keys,
                "selected_model_key": decision.selected_model_key,
                "confidence_tier": decision.confidence_tier,
                "budget_ratio": decision.budget_ratio,
                "canary_model_key": decision.canary_model_key,
                "used_canary": decision.used_canary,
                "routed_for_budget": decision.routed_for_budget,
            },
            "leaderboard": [score.__dict__ for score in scores if score.model_key in model_keys],
            "budget": self.budget_status(workspace_id=workspace_id, task_type=task_type),
            "active_canaries": [
                item.__dict__
                for item in self.list_canaries(
                    status="active",
                    workspace_id=workspace_id,
                    task_type=task_type,
                    user_segment=user_segment,
                )
            ],
        }

    def estimate_confidence(self, prompt: str, *, task_type: str | None = None) -> str:
        normalized = (prompt or "").strip().lower()
        if task_type in {"analysis", "architecture", "planning", "code", "strategy"}:
            return "high"
        complex_markers = ("analyze", "compare", "design", "debug", "refactor", "strategy", "explain why")
        if len(normalized) >= 600 or any(marker in normalized for marker in complex_markers):
            return "high"
        if len(normalized) >= 160 or normalized.count("\n") >= 4 or normalized.count("?") >= 2:
            return "medium"
        return "low"

    def sync_trace(self, trace_id: str, store: LearningStore) -> None:
        run = store.get_run_by_trace(trace_id)
        if run is None:
            return
        feedback = store.list_feedback(trace_id=trace_id)
        outcomes = store.list_outcomes(trace_id=trace_id)
        quality_score = _quality_from_feedback(feedback, outcomes)
        with self._connect() as connection:
            connection.execute(
                f"""
                UPDATE {self._table('events')}
                SET quality_score = ?, synced_at = ?
                WHERE trace_id = ?
                """,
                (quality_score, utcnow(), trace_id),
            )

    def canary_outcome_delta(
        self,
        *,
        workspace_id: str | None = None,
        task_type: str | None = None,
    ) -> float:
        clauses = ["workspace_id = ?"]
        params: list[Any] = [_segment(workspace_id)]
        if task_type is not None:
            clauses.append("task_type = ?")
            params.append(_segment(task_type))
        where_sql = " AND ".join(clauses)
        with self._connect() as connection:
            row = connection.execute(
                f"""
                SELECT
                    AVG(CASE WHEN was_canary = 1 THEN quality_score END) AS canary_quality,
                    AVG(CASE WHEN was_canary = 0 THEN quality_score END) AS baseline_quality
                FROM {self._table('events')}
                WHERE {where_sql}
                """,
                tuple(params),
            ).fetchone()
        if row is None:
            return 0.0
        canary_quality = float(row["canary_quality"] or 0.0)
        baseline_quality = float(row["baseline_quality"] or 0.0)
        return round(canary_quality - baseline_quality, 6)

    def _active_canary(
        self,
        *,
        workspace_id: str | None,
        task_type: str | None,
        user_segment: str | None,
        candidate_keys: list[str],
    ) -> CanaryDeployment | None:
        with self._connect() as connection:
            rows = connection.execute(
                f"""
                SELECT * FROM {self._table('canaries')}
                WHERE status = 'active'
                  AND workspace_id IN (?, '')
                  AND task_type IN (?, '')
                  AND user_segment IN (?, '')
                ORDER BY created_at DESC
                """,
                (_segment(workspace_id), _segment(task_type), _segment(user_segment)),
            ).fetchall()
        for row in rows:
            model_key = str(row["model_key"])
            if model_key not in candidate_keys:
                continue
            return CanaryDeployment(
                id=str(row["id"]),
                model_key=model_key,
                rollout_percentage=int(row["rollout_percentage"]),
                workspace_id=workspace_id,
                task_type=task_type,
                user_segment=user_segment,
                status=str(row["status"]),
                created_at=str(row["created_at"]),
                ended_at=str(row["ended_at"]) if row["ended_at"] is not None else None,
            )
        return None

    def _should_assign_canary(self, deployment: CanaryDeployment, assignment_key: str | None) -> bool:
        key = assignment_key or deployment.id
        digest = hashlib.sha256(f"{deployment.id}:{key}".encode()).hexdigest()
        bucket = int(digest[:8], 16) % 100
        return bucket < deployment.rollout_percentage

    def _spent_cost(self, *, workspace_id: str | None, task_type: str | None) -> float:
        clauses = ["workspace_id = ?"]
        params: list[Any] = [_segment(workspace_id)]
        if task_type is not None:
            clauses.append("task_type = ?")
            params.append(_segment(task_type))
        query = f"SELECT COALESCE(SUM(cost_usd), 0) AS total_cost FROM {self._table('events')} WHERE {' AND '.join(clauses)}"
        with self._connect() as connection:
            row = connection.execute(query, tuple(params)).fetchone()
        return float(row["total_cost"] if row is not None else 0.0)

    def _estimated_cost_for_model(self, model_key: str) -> float:
        _, _, model_name = model_key.partition(":")
        candidate = model_name or model_key
        for pattern, (input_rate, output_rate) in DEFAULT_RATES.items():
            if pattern == candidate or ("*" in pattern and fnmatch(candidate, pattern)):
                return (input_rate + output_rate) / 2_000_000
        return 0.0001


class _PostgreSQLCompatConnection:
    """Small adapter to reuse SQLite-shaped query code with psycopg."""

    def __init__(self, connection: Any) -> None:
        self._connection = connection

    def execute(self, query: str, params: tuple[Any, ...] | list[Any] | None = None) -> Any:
        return self._connection.execute(_translate_postgresql_query(query), tuple(params or ()))

    def cursor(self) -> Any:
        return self._connection.cursor()

    def commit(self) -> None:
        self._connection.commit()

    def close(self) -> None:
        self._connection.close()

    def __enter__(self) -> _PostgreSQLCompatConnection:
        self._connection.__enter__()
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        return bool(self._connection.__exit__(exc_type, exc, tb))


def _translate_postgresql_query(query: str) -> str:
    return query.replace("?", "%s")


__all__ = [
    "MODEL_ROUTING_SCHEMA_KIND",
    "MODEL_ROUTING_SCHEMA_VERSION",
    "CanaryDeployment",
    "ModelRoutingLearningEngine",
    "ModelScore",
    "RoutingBudget",
    "RoutingDecision",
]
