"""
Database migration script to add new handoff and analytics fields.
Run this after activating your virtual environment:
    python add_handoff_fields.py
"""

from sqlalchemy import create_engine, Column, Integer, String, DateTime, Boolean, Float, JSON, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import os
from dotenv import load_dotenv

load_dotenv()

# Database connection
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./customeragent.db")
engine = create_engine(DATABASE_URL)
Base = declarative_base()

def add_columns():
    """Add new columns to existing tables"""
    from sqlalchemy import text
    
    with engine.connect() as conn:
        # Add to chat_sessions
        try:
            conn.execute(text("""
                ALTER TABLE chat_sessions 
                ADD COLUMN urgency VARCHAR DEFAULT 'low'
            """))
            print("✓ Added urgency to chat_sessions")
        except Exception as e:
            print(f"  urgency column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE chat_sessions 
                ADD COLUMN agent_id INTEGER
            """))
            print("✓ Added agent_id to chat_sessions")
        except Exception as e:
            print(f"  agent_id column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE chat_sessions 
                ADD COLUMN takeover_at TIMESTAMP
            """))
            print("✓ Added takeover_at to chat_sessions")
        except Exception as e:
            print(f"  takeover_at column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE chat_sessions 
                ADD COLUMN lead_score INTEGER DEFAULT 0
            """))
            print("✓ Added lead_score to chat_sessions")
        except Exception as e:
            print(f"  lead_score column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE chat_sessions 
                ADD COLUMN qualification_data JSON
            """))
            print("✓ Added qualification_data to chat_sessions")
        except Exception as e:
            print(f"  qualification_data column may already exist: {e}")
        
        # Add to chat_messages
        try:
            conn.execute(text("""
                ALTER TABLE chat_messages 
                ADD COLUMN analysis JSON
            """))
            print("✓ Added analysis to chat_messages")
        except Exception as e:
            print(f"  analysis column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE chat_messages 
                ADD COLUMN confidence_score FLOAT
            """))
            print("✓ Added confidence_score to chat_messages")
        except Exception as e:
            print(f"  confidence_score column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE chat_messages 
                ADD COLUMN escalated BOOLEAN DEFAULT FALSE
            """))
            print("✓ Added escalated to chat_messages")
        except Exception as e:
            print(f"  escalated column may already exist: {e}")
        
        # Add to unanswered_questions
        try:
            conn.execute(text("""
                ALTER TABLE unanswered_questions 
                ADD COLUMN sla_deadline TIMESTAMP
            """))
            print("✓ Added sla_deadline to unanswered_questions")
        except Exception as e:
            print(f"  sla_deadline column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE unanswered_questions 
                ADD COLUMN sla_breached BOOLEAN DEFAULT FALSE
            """))
            print("✓ Added sla_breached to unanswered_questions")
        except Exception as e:
            print(f"  sla_breached column may already exist: {e}")
        
        try:
            conn.execute(text("""
                ALTER TABLE unanswered_questions 
                ADD COLUMN priority_score INTEGER DEFAULT 50
            """))
            print("✓ Added priority_score to unanswered_questions")
        except Exception as e:
            print(f"  priority_score column may already exist: {e}")
        
        conn.commit()
        print("\n✅ Migration completed successfully!")

if __name__ == "__main__":
    print("Starting database migration...")
    print(f"Database: {DATABASE_URL}\n")
    add_columns()