"""PostgreSQL Database Persistence Manager for GCP Solution Architecture Agent.

Manages connection pooling, schema initialization, and database persistence
for ADK sessions, workflow executions, artifacts, evaluations, and orchestrator review logs.
Supports external PostgreSQL database with SQLite fallback for offline development/testing.
"""

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

try:
    import psycopg2
    from psycopg2.extras import RealDictCursor
    HAS_PSYCOPG2 = True
except ImportError:
    HAS_PSYCOPG2 = False

from app.config import get_settings

logger = logging.getLogger(__name__)


class DatabaseManager:
    """Handles persistence of sessions, execution runs, artifacts, and review logs."""

    def __init__(self, db_url: Optional[str] = None) -> None:
        self.settings = get_settings()
        self.db_url = db_url or self.settings.get_database_url()
        self.is_postgres = self.db_url.startswith("postgresql://") or self.db_url.startswith("postgres://")
        self._init_db()

    def _get_connection(self):
        """Get database connection (PostgreSQL if available/accessible, else SQLite)."""
        if self.is_postgres and HAS_PSYCOPG2:
            try:
                conn = psycopg2.connect(self.db_url, connect_timeout=3)
                return conn
            except Exception as exc:
                logger.warning(
                    "PostgreSQL connection to %s failed (%s); falling back to local SQLite database %s",
                    self.db_url,
                    exc,
                    self.settings.SQLITE_FALLBACK_DB,
                )

        # SQLite fallback connection
        conn = sqlite3.connect(self.settings.SQLITE_FALLBACK_DB)
        conn.row_factory = sqlite3.Row
        return conn

    def _init_db(self) -> None:
        """Initialize database schema tables if they do not exist."""
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            
            # 1. Sessions Table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS adk_sessions (
                    session_id VARCHAR(128) PRIMARY KEY,
                    agent_name VARCHAR(128) NOT NULL,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                    state_data TEXT NOT NULL,
                    metadata TEXT
                )
            """)

            # 2. Workflow Executions Table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS workflow_executions (
                    execution_id VARCHAR(128) PRIMARY KEY,
                    session_id VARCHAR(128) NOT NULL,
                    status VARCHAR(32) NOT NULL,
                    current_phase VARCHAR(64) NOT NULL,
                    loop_count INTEGER DEFAULT 0,
                    started_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                    completed_at TIMESTAMP WITH TIME ZONE,
                    request_summary TEXT,
                    results_json TEXT
                )
            """)

            # 3. Artifacts Table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS adk_artifacts (
                    artifact_id VARCHAR(128) PRIMARY KEY,
                    execution_id VARCHAR(128) NOT NULL,
                    artifact_type VARCHAR(64) NOT NULL,
                    file_path VARCHAR(256) NOT NULL,
                    content TEXT NOT NULL,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # 4. Review & Orchestration Loop Logs Table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS orchestrator_review_logs (
                    log_id VARCHAR(128) PRIMARY KEY,
                    execution_id VARCHAR(128) NOT NULL,
                    iteration INTEGER NOT NULL,
                    review_status VARCHAR(32) NOT NULL,
                    reviewer_agent VARCHAR(64) NOT NULL,
                    feedback TEXT,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # 5. Evaluations Table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS adk_evaluations (
                    eval_id VARCHAR(128) PRIMARY KEY,
                    case_id VARCHAR(128) NOT NULL,
                    total_score FLOAT NOT NULL,
                    max_score FLOAT NOT NULL,
                    pass_rate FLOAT NOT NULL,
                    passed BOOLEAN NOT NULL,
                    details_json TEXT NOT NULL,
                    evaluated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
                )
            """)

            conn.commit()
            logger.info("Database schema initialized successfully.")
        except Exception as exc:
            logger.error("Failed to initialize database schema: %s", exc)
        finally:
            conn.close()

    def save_session(self, session_id: str, agent_name: str, state_data: Dict[str, Any], metadata: Optional[Dict[str, Any]] = None) -> bool:
        """Persist or update ADK session state in PostgreSQL."""
        conn = self._get_connection()
        now = datetime.now(timezone.utc).isoformat()
        try:
            cursor = conn.cursor()
            state_json = json.dumps(state_data)
            meta_json = json.dumps(metadata or {})

            if isinstance(conn, sqlite3.Connection):
                cursor.execute("""
                    INSERT INTO adk_sessions (session_id, agent_name, created_at, updated_at, state_data, metadata)
                    VALUES (?, ?, ?, ?, ?, ?)
                    ON CONFLICT(session_id) DO UPDATE SET
                        updated_at = excluded.updated_at,
                        state_data = excluded.state_data,
                        metadata = excluded.metadata
                """, (session_id, agent_name, now, now, state_json, meta_json))
            else:
                cursor.execute("""
                    INSERT INTO adk_sessions (session_id, agent_name, created_at, updated_at, state_data, metadata)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT(session_id) DO UPDATE SET
                        updated_at = EXCLUDED.updated_at,
                        state_data = EXCLUDED.state_data,
                        metadata = EXCLUDED.metadata
                """, (session_id, agent_name, now, now, state_json, meta_json))

            conn.commit()
            return True
        except Exception as exc:
            logger.error("Error saving session %s: %s", session_id, exc)
            return False
        finally:
            conn.close()

    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve ADK session state by session_id."""
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            if isinstance(conn, sqlite3.Connection):
                cursor.execute("SELECT * FROM adk_sessions WHERE session_id = ?", (session_id,))
            else:
                cursor.execute("SELECT * FROM adk_sessions WHERE session_id = %s", (session_id,))

            row = cursor.fetchone()
            if not row:
                return None

            if isinstance(conn, sqlite3.Connection):
                return {
                    "session_id": row["session_id"],
                    "agent_name": row["agent_name"],
                    "state_data": json.loads(row["state_data"]),
                    "metadata": json.loads(row["metadata"] or "{}"),
                    "updated_at": row["updated_at"],
                }
            else:
                return {
                    "session_id": row[0],
                    "agent_name": row[1],
                    "created_at": str(row[2]),
                    "updated_at": str(row[3]),
                    "state_data": json.loads(row[4]),
                    "metadata": json.loads(row[5] or "{}"),
                }
        except Exception as exc:
            logger.error("Error reading session %s: %s", session_id, exc)
            return None
        finally:
            conn.close()

    def save_workflow_execution(
        self,
        execution_id: str,
        session_id: str,
        status: str,
        current_phase: str,
        loop_count: int,
        request_summary: str,
        results: Dict[str, Any],
    ) -> bool:
        """Save workflow execution record to database."""
        conn = self._get_connection()
        now = datetime.now(timezone.utc).isoformat()
        try:
            cursor = conn.cursor()
            results_json = json.dumps(results)

            if isinstance(conn, sqlite3.Connection):
                cursor.execute("""
                    INSERT INTO workflow_executions
                    (execution_id, session_id, status, current_phase, loop_count, started_at, completed_at, request_summary, results_json)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT(execution_id) DO UPDATE SET
                        status = excluded.status,
                        current_phase = excluded.current_phase,
                        loop_count = excluded.loop_count,
                        completed_at = excluded.completed_at,
                        results_json = excluded.results_json
                """, (execution_id, session_id, status, current_phase, loop_count, now, now, request_summary, results_json))
            else:
                cursor.execute("""
                    INSERT INTO workflow_executions
                    (execution_id, session_id, status, current_phase, loop_count, started_at, completed_at, request_summary, results_json)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT(execution_id) DO UPDATE SET
                        status = EXCLUDED.status,
                        current_phase = EXCLUDED.current_phase,
                        loop_count = EXCLUDED.loop_count,
                        completed_at = EXCLUDED.completed_at,
                        results_json = EXCLUDED.results_json
                """, (execution_id, session_id, status, current_phase, loop_count, now, now, request_summary, results_json))

            conn.commit()
            return True
        except Exception as exc:
            logger.error("Error saving workflow execution %s: %s", execution_id, exc)
            return False
        finally:
            conn.close()

    def save_artifact(self, artifact_id: str, execution_id: str, artifact_type: str, file_path: str, content: str) -> bool:
        """Write ADK artifact metadata & content to database."""
        conn = self._get_connection()
        now = datetime.now(timezone.utc).isoformat()
        try:
            cursor = conn.cursor()
            if isinstance(conn, sqlite3.Connection):
                cursor.execute("""
                    INSERT INTO adk_artifacts (artifact_id, execution_id, artifact_type, file_path, content, created_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                    ON CONFLICT(artifact_id) DO UPDATE SET
                        content = excluded.content,
                        file_path = excluded.file_path
                """, (artifact_id, execution_id, artifact_type, file_path, content, now))
            else:
                cursor.execute("""
                    INSERT INTO adk_artifacts (artifact_id, execution_id, artifact_type, file_path, content, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT(artifact_id) DO UPDATE SET
                        content = EXCLUDED.content,
                        file_path = EXCLUDED.file_path
                """, (artifact_id, execution_id, artifact_type, file_path, content, now))

            conn.commit()
            return True
        except Exception as exc:
            logger.error("Error saving artifact %s: %s", artifact_id, exc)
            return False
        finally:
            conn.close()

    def record_orchestrator_log(
        self,
        log_id: str,
        execution_id: str,
        iteration: int,
        review_status: str,
        reviewer_agent: str,
        feedback: str,
    ) -> bool:
        """Record orchestrator loop iteration review event into PostgreSQL log table."""
        conn = self._get_connection()
        now = datetime.now(timezone.utc).isoformat()
        try:
            cursor = conn.cursor()
            if isinstance(conn, sqlite3.Connection):
                cursor.execute("""
                    INSERT INTO orchestrator_review_logs (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, created_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                """, (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, now))
            else:
                cursor.execute("""
                    INSERT INTO orchestrator_review_logs (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s, %s)
                """, (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, now))

            conn.commit()
            return True
        except Exception as exc:
            logger.error("Error recording orchestrator review log %s: %s", log_id, exc)
            return False
        finally:
            conn.close()

    def save_evaluation(self, eval_id: str, case_id: str, total_score: float, max_score: float, pass_rate: float, passed: bool, details: Dict[str, Any]) -> bool:
        """Save ADK evaluation benchmark result to database."""
        conn = self._get_connection()
        now = datetime.now(timezone.utc).isoformat()
        try:
            cursor = conn.cursor()
            details_json = json.dumps(details)
            if isinstance(conn, sqlite3.Connection):
                cursor.execute("""
                    INSERT INTO adk_evaluations (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, evaluated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """, (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, now))
            else:
                cursor.execute("""
                    INSERT INTO adk_evaluations (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, evaluated_at)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                """, (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, now))

            conn.commit()
            return True
        except Exception as exc:
            logger.error("Error saving evaluation %s: %s", eval_id, exc)
            return False
        finally:
            conn.close()


_db_manager: Optional[DatabaseManager] = None


def get_db_manager() -> DatabaseManager:
    """Get singleton DatabaseManager instance."""
    global _db_manager
    if _db_manager is None:
        _db_manager = DatabaseManager()
    return _db_manager