import os
import psycopg2
from psycopg2.extras import RealDictCursor
from typing import List, Dict
import logging

logger = logging.getLogger(__name__)


class ChatDatabaseService:
    """Simple service for managing chat history in PostgreSQL"""

    def __init__(self):
        self.conn = psycopg2.connect(
            host=os.getenv("POSTGRES_HOST", "postgres"),
            user=os.getenv("POSTGRES_USER"),
            password=os.getenv("POSTGRES_PASSWORD"),
            dbname=os.getenv("POSTGRES_DB"),
        )

    def get_chat_history(self, student_id: str, limit: int = 20) -> List[Dict[str, str]]:
        """Get chat history for a student, returns in chronological order"""
        with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                """
                SELECT role, content
                FROM chat_history
                WHERE student_id = %s
                ORDER BY created_at DESC
                LIMIT %s;
                """,
                (student_id, limit)
            )
            results = cur.fetchall()
            # Return in chronological order (oldest first)
            return [{"role": row["role"], "content": row["content"]} for row in reversed(results)]

    def add_message(self, student_id: str, role: str, content: str):
        """Add a message to chat history"""
        with self.conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO chat_history (student_id, role, content)
                VALUES (%s, %s, %s);
                """,
                (student_id, role, content)
            )
            self.conn.commit()

    def clear_history(self, student_id: str):
        """Clear chat history for a student"""
        with self.conn.cursor() as cur:
            cur.execute(
                "DELETE FROM chat_history WHERE student_id = %s",
                (student_id,)
            )
            self.conn.commit()

    def limit_history(self, student_id: str, max_messages: int = 40):
        """Keep only recent messages for a student"""
        with self.conn.cursor() as cur:
            cur.execute(
                """
                DELETE FROM chat_history 
                WHERE student_id = %s 
                AND role != 'system'
                AND id NOT IN (
                    SELECT id FROM chat_history 
                    WHERE student_id = %s AND role != 'system'
                    ORDER BY created_at DESC 
                    LIMIT %s
                );
                """,
                (student_id, student_id, max_messages)
            )
            self.conn.commit()

    def close(self):
        if self.conn:
            self.conn.close()