import os
import psycopg2
import uuid
import datetime
import logging
import json

logger = logging.getLogger(__name__)

class PostgresMemoryManager:
    def __init__(self):
        self.db_url = os.getenv("NEON_DATABASE_URL")
        self.conn = None
        self._connect()

    def _connect(self):
        if not self.db_url:
            logger.error("No NEON_DATABASE_URL provided. Cannot connect to DB.")
            return

        try:
            self.conn = psycopg2.connect(self.db_url)
            self._initialize_tables()
            logger.info("Connected to PostgreSQL DB.")
        except Exception as e:
            logger.error(f"Failed to connect to database: {e}.")
            self.conn = None

    def _ensure_connection(self):
        if self.db_url:
            if self.conn is None or self.conn.closed != 0:
                self._connect()

    def _initialize_tables(self):
        if not self.conn:
            return
        try:
            with self.conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS sessions (
                        id VARCHAR(255) PRIMARY KEY,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        metadata TEXT
                    );
                """)
                cur.execute("""
                    ALTER TABLE sessions ADD COLUMN IF NOT EXISTS ip_address TEXT;
                    ALTER TABLE sessions ADD COLUMN IF NOT EXISTS location_data TEXT;
                """)
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS conversations (
                        id VARCHAR(255) PRIMARY KEY,
                        session_id VARCHAR(255) REFERENCES sessions(id),
                        user_message TEXT,
                        ai_response TEXT,
                        timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    );
                """)
            self.conn.commit()
        except Exception as e:
            logger.error(f"Error initializing tables: {e}")
            try:
                self.conn.rollback()
            except Exception:
                pass

    def create_session(self, metadata="", ip_address=None, location_data=None):
        self._ensure_connection()
        session_id = str(uuid.uuid4())
        location_data_str = json.dumps(location_data) if location_data else None

        if self.conn:
            try:
                with self.conn.cursor() as cur:
                    cur.execute(
                        "INSERT INTO sessions (id, metadata, ip_address, location_data) VALUES (%s, %s, %s, %s)",
                        (session_id, metadata, ip_address, location_data_str)
                    )
                self.conn.commit()
            except Exception as e:
                logger.error(f"DB Error creating session: {e}")
                try:
                    self.conn.rollback()
                except Exception:
                    self.conn = None
        return session_id

    def add_message(self, session_id, user_message, ai_response):
        self._ensure_connection()
        conv_id = str(uuid.uuid4())
        if self.conn:
            try:
                with self.conn.cursor() as cur:
                    cur.execute(
                        "INSERT INTO conversations (id, session_id, user_message, ai_response) VALUES (%s, %s, %s, %s)",
                        (conv_id, session_id, user_message, ai_response)
                    )
                self.conn.commit()
            except Exception as e:
                logger.error(f"DB Error adding message: {e}")
                try:
                    self.conn.rollback()
                except Exception:
                    self.conn = None

    def get_history(self, session_id):
        self._ensure_connection()
        if self.conn:
            try:
                with self.conn.cursor() as cur:
                    cur.execute(
                        "SELECT user_message, ai_response, timestamp FROM conversations WHERE session_id = %s ORDER BY timestamp ASC",
                        (session_id,)
                    )
                    rows = cur.fetchall()
                    return [{"user_message": row[0], "ai_response": row[1], "timestamp": row[2].isoformat()} for row in rows]
            except Exception as e:
                logger.error(f"DB Error getting history: {e}")
                try:
                    self.conn.rollback()
                except Exception:
                    self.conn = None
        
        return []