"""
PostgreSQL backend implementation for state storage.
Better for complex queries, large states, and analytics.
"""

import json
from typing import Optional, List
from datetime import datetime
import psycopg2
from psycopg2.extras import RealDictCursor
from psycopg2.pool import SimpleConnectionPool

from .base import StateBackend
from ..schemas import UnifiedState


class PostgreSQLBackend(StateBackend):
    """
    PostgreSQL-based state storage backend.
    
    Better for:
    - Complex queries and analytics
    - Large state sizes
    - ACID transactions
    - Long-term persistence
    """
    
    def __init__(
        self,
        host: str = "localhost",
        port: int = 5432,
        database: str = "uass",
        user: str = "postgres",
        password: Optional[str] = None,
        min_conn: int = 1,
        max_conn: int = 10,
    ):
        """
        Initialize PostgreSQL backend.
        
        Args:
            host: PostgreSQL host
            port: PostgreSQL port
            database: Database name
            user: Database user
            password: Database password
            min_conn: Minimum connection pool size
            max_conn: Maximum connection pool size
        """
        self.conn_params = {
            "host": host,
            "port": port,
            "database": database,
            "user": user,
            "password": password,
        }
        
        # Create connection pool
        self.pool = SimpleConnectionPool(
            min_conn, max_conn, **self.conn_params
        )
        
        # Initialize schema
        self._init_schema()
    
    def _get_conn(self):
        """Get connection from pool."""
        return self.pool.getconn()
    
    def _put_conn(self, conn):
        """Return connection to pool."""
        self.pool.putconn(conn)
    
    def _init_schema(self):
        """Initialize database schema."""
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                # Create states table
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS uass_states (
                        session_id VARCHAR(255) NOT NULL,
                        agent_id VARCHAR(255) NOT NULL,
                        state_data JSONB NOT NULL,
                        metadata JSONB NOT NULL,
                        created_at TIMESTAMP NOT NULL,
                        updated_at TIMESTAMP NOT NULL,
                        PRIMARY KEY (session_id, agent_id)
                    );
                """)
                
                # Create shared states table
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS uass_shared_states (
                        session_id VARCHAR(255) PRIMARY KEY,
                        state_data JSONB NOT NULL,
                        metadata JSONB NOT NULL,
                        created_at TIMESTAMP NOT NULL,
                        updated_at TIMESTAMP NOT NULL
                    );
                """)
                
                # Create state history table
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS uass_state_history (
                        id SERIAL PRIMARY KEY,
                        session_id VARCHAR(255) NOT NULL,
                        agent_id VARCHAR(255) NOT NULL,
                        state_data JSONB NOT NULL,
                        metadata JSONB NOT NULL,
                        change_type VARCHAR(50) NOT NULL,
                        timestamp TIMESTAMP NOT NULL DEFAULT NOW()
                    );
                """)
                
                # Create indexes
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_states_session 
                    ON uass_states(session_id);
                """)
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_history_session_agent 
                    ON uass_state_history(session_id, agent_id);
                """)
                cur.execute("""
                    CREATE INDEX IF NOT EXISTS idx_history_timestamp 
                    ON uass_state_history(timestamp);
                """)
                
                conn.commit()
        except Exception as e:
            conn.rollback()
            raise ConnectionError(f"Failed to initialize schema: {e}")
        finally:
            self._put_conn(conn)
    
    def write_state(self, state: UnifiedState) -> bool:
        """Write state to PostgreSQL."""
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                # Upsert state
                cur.execute("""
                    INSERT INTO uass_states 
                    (session_id, agent_id, state_data, metadata, created_at, updated_at)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT (session_id, agent_id)
                    DO UPDATE SET
                        state_data = EXCLUDED.state_data,
                        metadata = EXCLUDED.metadata,
                        updated_at = EXCLUDED.updated_at;
                """, (
                    state.session_id,
                    state.agent_id,
                    json.dumps(state.state.to_dict()),
                    json.dumps(state.metadata.to_dict()),
                    state.metadata.created_at,
                    state.metadata.updated_at,
                ))
                
                # Insert history
                cur.execute("""
                    INSERT INTO uass_state_history
                    (session_id, agent_id, state_data, metadata, change_type, timestamp)
                    VALUES (%s, %s, %s, %s, %s, %s);
                """, (
                    state.session_id,
                    state.agent_id,
                    json.dumps(state.state.to_dict()),
                    json.dumps(state.metadata.to_dict()),
                    "update",
                    state.metadata.updated_at,
                ))
                
                conn.commit()
                return True
        except Exception as e:
            conn.rollback()
            print(f"Error writing state to PostgreSQL: {e}")
            return False
        finally:
            self._put_conn(conn)
    
    def read_state(self, session_id: str, agent_id: str) -> Optional[UnifiedState]:
        """Read state from PostgreSQL."""
        conn = self._get_conn()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute("""
                    SELECT state_data, metadata, created_at, updated_at
                    FROM uass_states
                    WHERE session_id = %s AND agent_id = %s;
                """, (session_id, agent_id))
                
                row = cur.fetchone()
                if row is None:
                    return None
                
                state_data = row["state_data"]
                metadata_dict = row["metadata"]
                
                # Reconstruct UnifiedState
                from ..schemas import AgentState, StateMetadata
                
                agent_state = AgentState.from_dict(state_data)
                metadata = StateMetadata(
                    created_at=row["created_at"],
                    updated_at=row["updated_at"],
                    version=metadata_dict.get("version", 1),
                    framework=metadata_dict.get("framework", ""),
                    framework_version=metadata_dict.get("framework_version"),
                )
                
                return UnifiedState(
                    session_id=session_id,
                    agent_id=agent_id,
                    state=agent_state,
                    metadata=metadata,
                )
        except Exception as e:
            print(f"Error reading state from PostgreSQL: {e}")
            return None
        finally:
            self._put_conn(conn)
    
    def delete_state(self, session_id: str, agent_id: str) -> bool:
        """Delete state from PostgreSQL."""
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    DELETE FROM uass_states
                    WHERE session_id = %s AND agent_id = %s;
                """, (session_id, agent_id))
                conn.commit()
                return True
        except Exception as e:
            conn.rollback()
            print(f"Error deleting state from PostgreSQL: {e}")
            return False
        finally:
            self._put_conn(conn)
    
    def list_agents(self, session_id: str) -> List[str]:
        """List all agent IDs in a session."""
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    SELECT DISTINCT agent_id
                    FROM uass_states
                    WHERE session_id = %s;
                """, (session_id,))
                return [row[0] for row in cur.fetchall()]
        except Exception as e:
            print(f"Error listing agents from PostgreSQL: {e}")
            return []
        finally:
            self._put_conn(conn)
    
    def read_shared_state(self, session_id: str) -> Optional[UnifiedState]:
        """Read shared state for a session."""
        conn = self._get_conn()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute("""
                    SELECT state_data, metadata, created_at, updated_at
                    FROM uass_shared_states
                    WHERE session_id = %s;
                """, (session_id,))
                
                row = cur.fetchone()
                if row is None:
                    return None
                
                state_data = row["state_data"]
                metadata_dict = row["metadata"]
                
                from ..schemas import AgentState, StateMetadata
                
                agent_state = AgentState.from_dict(state_data)
                metadata = StateMetadata(
                    created_at=row["created_at"],
                    updated_at=row["updated_at"],
                    version=metadata_dict.get("version", 1),
                    framework=metadata_dict.get("framework", "shared"),
                    framework_version=metadata_dict.get("framework_version"),
                )
                
                return UnifiedState(
                    session_id=session_id,
                    agent_id="shared",
                    state=agent_state,
                    metadata=metadata,
                )
        except Exception as e:
            print(f"Error reading shared state from PostgreSQL: {e}")
            return None
        finally:
            self._put_conn(conn)
    
    def write_shared_state(self, session_id: str, state: UnifiedState) -> bool:
        """Write shared state for a session."""
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO uass_shared_states
                    (session_id, state_data, metadata, created_at, updated_at)
                    VALUES (%s, %s, %s, %s, %s)
                    ON CONFLICT (session_id)
                    DO UPDATE SET
                        state_data = EXCLUDED.state_data,
                        metadata = EXCLUDED.metadata,
                        updated_at = EXCLUDED.updated_at;
                """, (
                    session_id,
                    json.dumps(state.state.to_dict()),
                    json.dumps(state.metadata.to_dict()),
                    state.metadata.created_at,
                    state.metadata.updated_at,
                ))
                conn.commit()
                return True
        except Exception as e:
            conn.rollback()
            print(f"Error writing shared state to PostgreSQL: {e}")
            return False
        finally:
            self._put_conn(conn)
    
    def query_states(self, query: str, params: tuple = ()) -> List[dict]:
        """
        Execute custom SQL query on states.
        PostgreSQL-specific feature for complex queries.
        
        Example:
            backend.query_states(
                "SELECT * FROM uass_states WHERE metadata->>'framework' = %s",
                ('langgraph',)
            )
        """
        conn = self._get_conn()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(query, params)
                return [dict(row) for row in cur.fetchall()]
        except Exception as e:
            print(f"Error executing query: {e}")
            return []
        finally:
            self._put_conn(conn)

