"""Reflection memory for lessons learned."""

from __future__ import annotations

import json
import os
import sqlite3
from pathlib import Path

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 IMPROVEMENT_SCHEMA_KIND, IMPROVEMENT_SCHEMA_VERSION, ReflectionLesson


class ReflectionMemory:
    """Persistent structured memory for lessons learned."""

    def __init__(
        self,
        path: str | None = None,
        table_name: str = "reflection_lessons",
        dsn: str | None = None,
    ) -> None:
        """Initialize reflection memory.

        Args:
            path: SQLite database path
            table_name: Table name for lessons
            dsn: PostgreSQL DSN (optional)
        """
        self.path = path or os.getenv(
            "LOGOS_ADK_REFLECTION_DB_PATH", ".logos_adk/reflection_memory.db"
        )
        self.dsn = resolve_database_url(
            dsn or os.getenv("LOGOS_ADK_REFLECTION_DSN"),
            env_vars=("LOGOS_ADK_REFLECTION_DSN", "LOGOS_ADK_DATABASE_URL", "DATABASE_URL"),
        )
        self.table_name = table_name

    def _connect(self):
        """Create database connection."""
        if self.dsn:
            connection = import_psycopg().connect(self.dsn)
            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) -> None:
        """Ensure database schema is up to date."""
        if self.dsn:
            self._ensure_postgresql_schema(connection)
        else:
            self._ensure_sqlite_schema(connection)
        connection.commit()

    def _ensure_postgresql_schema(self, connection) -> None:
        """Ensure PostgreSQL schema."""
        current = get_postgresql_schema_version(connection, IMPROVEMENT_SCHEMA_KIND) or 0
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id TEXT PRIMARY KEY,
                    workspace_id TEXT,
                    agent_name TEXT,
                    category TEXT NOT NULL,
                    lesson TEXT NOT NULL,
                    rationale TEXT NOT NULL,
                    suggested_action TEXT NOT NULL,
                    confidence DOUBLE PRECISION NOT NULL,
                    source_run_ids_json JSONB NOT NULL,
                    metadata_json JSONB NOT NULL,
                    created_at TEXT NOT NULL
                )
                """
            )
            cursor.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_lookup
                ON {self.table_name} (workspace_id, agent_name, category, created_at DESC)
                """
            )
        if current < IMPROVEMENT_SCHEMA_VERSION:
            set_postgresql_schema_version(
                connection, IMPROVEMENT_SCHEMA_KIND, IMPROVEMENT_SCHEMA_VERSION
            )

    def _ensure_sqlite_schema(self, connection) -> None:
        """Ensure SQLite schema."""
        current = get_sqlite_schema_version(connection, IMPROVEMENT_SCHEMA_KIND) or 0
        connection.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                id TEXT PRIMARY KEY,
                workspace_id TEXT,
                agent_name TEXT,
                category TEXT NOT NULL,
                lesson TEXT NOT NULL,
                rationale TEXT NOT NULL,
                suggested_action TEXT NOT NULL,
                confidence REAL NOT NULL,
                source_run_ids_json TEXT NOT NULL,
                metadata_json TEXT NOT NULL,
                created_at TEXT NOT NULL
            )
            """
        )
        connection.execute(
            f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_lookup ON {self.table_name} (workspace_id, agent_name, category, created_at DESC)"
        )
        if current < IMPROVEMENT_SCHEMA_VERSION:
            set_sqlite_schema_version(
                connection, IMPROVEMENT_SCHEMA_KIND, IMPROVEMENT_SCHEMA_VERSION
            )

    def save_lesson(self, lesson: ReflectionLesson) -> ReflectionLesson:
        """Save lesson to memory."""
        with self._connect() as connection:
            if self.dsn:
                with connection.cursor() as cursor:
                    cursor.execute(
                        f"""
                        INSERT INTO {self.table_name}
                        (id, workspace_id, agent_name, category, lesson, rationale, suggested_action,
                         confidence, source_run_ids_json, metadata_json, created_at)
                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s)
                        ON CONFLICT(id) DO UPDATE SET
                            workspace_id = excluded.workspace_id,
                            agent_name = excluded.agent_name,
                            category = excluded.category,
                            lesson = excluded.lesson,
                            rationale = excluded.rationale,
                            suggested_action = excluded.suggested_action,
                            confidence = excluded.confidence,
                            source_run_ids_json = excluded.source_run_ids_json,
                            metadata_json = excluded.metadata_json,
                            created_at = excluded.created_at
                        """,
                        (
                            lesson.id,
                            lesson.workspace_id,
                            lesson.agent_name,
                            lesson.category,
                            lesson.lesson,
                            lesson.rationale,
                            lesson.suggested_action,
                            lesson.confidence,
                            json.dumps(lesson.source_run_ids),
                            json.dumps(lesson.metadata),
                            lesson.created_at,
                        ),
                    )
            else:
                connection.execute(
                    f"""
                    INSERT OR REPLACE INTO {self.table_name}
                    (id, workspace_id, agent_name, category, lesson, rationale, suggested_action,
                     confidence, source_run_ids_json, metadata_json, created_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        lesson.id,
                        lesson.workspace_id,
                        lesson.agent_name,
                        lesson.category,
                        lesson.lesson,
                        lesson.rationale,
                        lesson.suggested_action,
                        lesson.confidence,
                        json.dumps(lesson.source_run_ids),
                        json.dumps(lesson.metadata),
                        lesson.created_at,
                    ),
                )
        return lesson

    def list_lessons(
        self,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        category: str | None = None,
        limit: int = 50,
    ) -> list[ReflectionLesson]:
        """List lessons with filters."""
        clauses = []
        params = []
        if workspace_id is not None:
            placeholders = "%s" if self.dsn else "?"
            clauses.append(f"workspace_id = {placeholders}")
            params.append(workspace_id)
        if agent_name is not None:
            placeholders = "%s" if self.dsn else "?"
            clauses.append(f"agent_name = {placeholders}")
            params.append(agent_name)
        if category is not None:
            placeholders = "%s" if self.dsn else "?"
            clauses.append(f"category = {placeholders}")
            params.append(category)
        where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""

        limit_placeholder = "%s" if self.dsn else "?"
        params.append(limit)

        with self._connect() as connection:
            if self.dsn:
                with connection.cursor() as cursor:
                    cursor.execute(
                        f"""
                        SELECT * FROM {self.table_name}
                        {where_sql}
                        ORDER BY created_at DESC
                        LIMIT {limit_placeholder}
                        """,
                        params,
                    )
                    rows = cursor.fetchall()
            else:
                rows = connection.execute(
                    f"""
                    SELECT * FROM {self.table_name}
                    {where_sql}
                    ORDER BY created_at DESC
                    LIMIT {limit_placeholder}
                    """,
                    params,
                ).fetchall()

        return [self._row_to_lesson(row) for row in rows]

    def get_lesson(self, lesson_id: str) -> ReflectionLesson | None:
        """Get lesson by ID."""
        with self._connect() as connection:
            row = connection.execute(
                f"SELECT * FROM {self.table_name} WHERE id = ?",
                (lesson_id,),
            ).fetchone()
        return self._row_to_lesson(row) if row else None

    def delete_lesson(self, lesson_id: str) -> bool:
        """Delete lesson by ID."""
        with self._connect() as connection:
            result = connection.execute(
                f"DELETE FROM {self.table_name} WHERE id = ?",
                (lesson_id,),
            )
            return result.rowcount > 0

    def _row_to_lesson(self, row) -> ReflectionLesson:
        """Convert database row to lesson."""
        if hasattr(row, "keys"):
            data = dict(row)
        else:
            data = {
                "id": row[0],
                "workspace_id": row[1],
                "agent_name": row[2],
                "category": row[3],
                "lesson": row[4],
                "rationale": row[5],
                "suggested_action": row[6],
                "confidence": row[7],
                "source_run_ids_json": row[8],
                "metadata_json": row[9],
                "created_at": row[10],
            }

        source_run_ids = data["source_run_ids_json"]
        if isinstance(source_run_ids, str):
            source_run_ids = json.loads(source_run_ids) if source_run_ids else []
        elif source_run_ids is None:
            source_run_ids = []

        metadata = data["metadata_json"]
        if isinstance(metadata, str):
            metadata = json.loads(metadata) if metadata else {}
        elif metadata is None:
            metadata = {}

        return ReflectionLesson(
            id=str(data["id"]),
            workspace_id=data["workspace_id"],
            agent_name=data["agent_name"],
            category=str(data["category"]),
            lesson=str(data["lesson"]),
            rationale=str(data["rationale"]),
            suggested_action=str(data["suggested_action"]),
            confidence=float(data["confidence"]),
            source_run_ids=source_run_ids,
            metadata=metadata,
            created_at=str(data["created_at"]),
        )

    def search_lessons(
        self,
        query: str,
        *,
        workspace_id: str | None = None,
        limit: int = 20,
    ) -> list[ReflectionLesson]:
        """Search lessons by text query."""
        search_pattern = f"%{query}%"
        with self._connect() as connection:
            rows = connection.execute(
                f"""
                SELECT * FROM {self.table_name}
                WHERE (workspace_id = ? OR ? IS NULL)
                  AND (lesson LIKE ? OR rationale LIKE ? OR suggested_action LIKE ?)
                ORDER BY confidence DESC, created_at DESC
                LIMIT ?
                """,
                (workspace_id, workspace_id, search_pattern, search_pattern, search_pattern, limit),
            ).fetchall()
        return [self._row_to_lesson(row) for row in rows]


__all__ = ["ReflectionMemory"]
