import logging
from app.db.session import engine
from sqlalchemy import text, inspect

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def run_migrations():
    """
    Checks for missing columns and adds them if necessary.
    This is a simple idempotent migration system to handle schema updates.
    """
    logger.info("Starting database migrations...")
    
    with engine.connect() as conn:
        inspector = inspect(engine)
        if not inspector.has_table("messages"):
            logger.info("Table 'messages' does not exist, skipping migrations (will be handled by Base.metadata.create_all).")
            return

        columns = [c["name"] for c in inspector.get_columns("messages")]
        
        # List of (column_name, column_type) to ensure existence
        required_columns = [
            ("audio_path", "TEXT"),
            ("model_response_time", "INTEGER"),
            ("token_count", "INTEGER"),
            ("reasoning_content", "TEXT")
        ]

        
        for col_name, col_type in required_columns:
            if col_name not in columns:
                logger.info(f"Adding column '{col_name}' to 'messages' table...")
                try:
                    conn.execute(text(f"ALTER TABLE messages ADD COLUMN {col_name} {col_type}"))
                    conn.commit()
                    logger.info(f"Successfully added '{col_name}'.")
                except Exception as e:
                    logger.error(f"Failed to add column '{col_name}': {e}")
            else:
                logger.info(f"Column '{col_name}' already exists in 'messages'.")

        # Session table migrations
        session_columns = [c["name"] for c in inspector.get_columns("sessions")]
        session_required_columns = [
            ("stt_provider_name",  "TEXT"),
            ("tts_provider_name",  "TEXT"),
            # M3: Agent Node attachment
            ("sync_workspace_id",  "TEXT"),
            ("attached_node_ids",  "TEXT"),
            ("node_sync_status",   "TEXT"),
            ("sync_config",        "TEXT"),
            ("is_cancelled",       "INTEGER DEFAULT 0"),
            ("restrict_skills",    "BOOLEAN DEFAULT 0"),
            ("allowed_skill_names","TEXT"),
            ("system_prompt_override","TEXT"),
            ("is_locked",          "BOOLEAN DEFAULT 0"),
            ("model_name",         "TEXT"),
        ]
        for col_name, col_type in session_required_columns:
            if col_name not in session_columns:
                logger.info(f"Adding column '{col_name}' to 'sessions' table...")
                try:
                    conn.execute(text(f"ALTER TABLE sessions ADD COLUMN {col_name} {col_type}"))
                    conn.commit()
                    logger.info(f"Successfully added '{col_name}'.")
                except Exception as e:
                    logger.error(f"Failed to add column '{col_name}': {e}")
            else:
                logger.info(f"Column '{col_name}' already exists in 'sessions'.")

        # Users table migrations for local authentication
        user_columns = [c["name"] for c in inspector.get_columns("users")]
        if "password_hash" not in user_columns:
            logger.info("Adding column 'password_hash' to 'users' table...")
            try:
                conn.execute(text("ALTER TABLE users ADD COLUMN password_hash TEXT"))
                conn.commit()
                logger.info("Successfully added 'password_hash'.")
            except Exception as e:
                logger.error(f"Failed to add column 'password_hash': {e}")

        # --- M6: Agent Node Tables ---
        # Create agent_nodes table if it doesn't exist
        if not inspector.has_table("agent_nodes"):
            logger.info("Creating table 'agent_nodes'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS agent_nodes (
                        id            INTEGER  PRIMARY KEY AUTOINCREMENT,
                        node_id       TEXT     UNIQUE NOT NULL,
                        display_name  TEXT     NOT NULL,
                        description   TEXT,
                        registered_by TEXT     NOT NULL,
                        skill_config  TEXT     NOT NULL DEFAULT '{}',
                        capabilities  TEXT     DEFAULT '{}',
                        invite_token  TEXT     UNIQUE,
                        is_active     INTEGER  NOT NULL DEFAULT 1,
                        last_status   TEXT     NOT NULL DEFAULT 'offline',
                        last_seen_at  DATETIME,
                        created_at    DATETIME DEFAULT CURRENT_TIMESTAMP
                    )
                """))
                conn.commit()
                logger.info("Table 'agent_nodes' created.")
            except Exception as e:
                logger.error(f"Failed to create 'agent_nodes': {e}")
        else:
            # Table exists — ensure all columns are present
            node_columns = [c["name"] for c in inspector.get_columns("agent_nodes")]
            node_required_columns = [
                ("display_name",   "TEXT"),
                ("registered_by",  "TEXT"),
                ("skill_config",   "TEXT"),
                ("invite_token",   "TEXT"),
                ("is_active",      "INTEGER"),
                ("last_status",    "TEXT"),
                ("last_seen_at",   "DATETIME"),
                ("capabilities",   "TEXT"),
            ]
            for col_name, col_type in node_required_columns:
                if col_name not in node_columns:
                    logger.info(f"Adding column '{col_name}' to 'agent_nodes' table...")
                    try:
                        conn.execute(text(f"ALTER TABLE agent_nodes ADD COLUMN {col_name} {col_type}"))
                        conn.commit()
                    except Exception as e:
                        logger.error(f"Failed to add column '{col_name}': {e}")

        # Create node_group_access table if it doesn't exist
        if not inspector.has_table("node_group_access"):
            logger.info("Creating table 'node_group_access'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS node_group_access (
                        id           INTEGER  PRIMARY KEY AUTOINCREMENT,
                        node_id      TEXT     NOT NULL REFERENCES agent_nodes(node_id),
                        group_id     TEXT     NOT NULL REFERENCES groups(id),
                        access_level TEXT     NOT NULL DEFAULT 'use',
                        granted_by   TEXT     NOT NULL REFERENCES users(id),
                        granted_at   DATETIME DEFAULT CURRENT_TIMESTAMP
                    )
                """))
                conn.commit()
                logger.info("Table 'node_group_access' created.")
            except Exception as e:
                logger.error(f"Failed to create 'node_group_access': {e}")

        # Create skill_group_access table if it doesn't exist
        if not inspector.has_table("skill_group_access"):
            logger.info("Creating table 'skill_group_access'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS skill_group_access (
                        id           INTEGER  PRIMARY KEY AUTOINCREMENT,
                        skill_id     INTEGER  NOT NULL REFERENCES skills(id),
                        group_id     TEXT     NOT NULL REFERENCES groups(id),
                        granted_by   TEXT     NOT NULL REFERENCES users(id),
                        granted_at   DATETIME DEFAULT CURRENT_TIMESTAMP
                    )
                """))
                conn.commit()
                logger.info("Table 'skill_group_access' created.")
            except Exception as e:
                logger.error(f"Failed to create 'skill_group_access': {e}")

        # Create session_skills table if it doesn't exist
        if not inspector.has_table("session_skills"):
            logger.info("Creating table 'session_skills'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS session_skills (
                        session_id INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
                        skill_id   INTEGER NOT NULL REFERENCES skills(id) ON DELETE CASCADE,
                        PRIMARY KEY (session_id, skill_id)
                    )
                """))
                conn.commit()
                logger.info("Table 'session_skills' created.")
            except Exception as e:
                logger.error(f"Failed to create 'session_skills': {e}")

        # Create skill_files table if it doesn't exist
        if not inspector.has_table("skill_files"):
            logger.info("Creating table 'skill_files'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS skill_files (
                        id         INTEGER PRIMARY KEY AUTOINCREMENT,
                        skill_id   INTEGER NOT NULL REFERENCES skills(id),
                        file_path  TEXT    NOT NULL,
                        content    TEXT,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                    )
                """))
                
                conn.execute(text("CREATE INDEX IF NOT EXISTS idx_skill_files_skill_id ON skill_files(skill_id)"))
                conn.execute(text("CREATE INDEX IF NOT EXISTS idx_skill_files_file_path ON skill_files(file_path)"))
                
                conn.commit()
                logger.info("Table 'skill_files' created.")
            except Exception as e:
                logger.error(f"Failed to create 'skill_files': {e}")

        if inspector.has_table("skills"):
            skill_columns = [c["name"] for c in inspector.get_columns("skills")]
            skill_required_columns = [
                ("is_enabled",    "INTEGER DEFAULT 1"),
                ("features",      "TEXT DEFAULT '[\"chat\"]'"),
                ("is_system",     "INTEGER DEFAULT 0"),
                ("skill_type",    "TEXT DEFAULT 'local'"),
                ("extra_metadata", "TEXT DEFAULT '{}'")
            ]
            for col_name, col_type in skill_required_columns:
                if col_name not in skill_columns:
                    logger.info(f"Adding column '{col_name}' to 'skills' table...")
                    try:
                        conn.execute(text(f"ALTER TABLE skills ADD COLUMN {col_name} {col_type}"))
                        conn.commit()
                        logger.info(f"Successfully added '{col_name}' to 'skills'.")
                    except Exception as e:
                        logger.error(f"Failed to add column '{col_name}' to 'skills': {e}")
                        
            # Remove deprecated columns
            deprecated_columns = ["system_prompt", "preview_markdown", "config"]
            for col_name in deprecated_columns:
                if col_name in skill_columns:
                    logger.info(f"Dropping deprecated column '{col_name}' from 'skills' table...")
                    try:
                        conn.execute(text(f"ALTER TABLE skills DROP COLUMN {col_name}"))
                        conn.commit()
                        logger.info(f"Successfully dropped '{col_name}' from 'skills'.")
                    except Exception as e:
                        logger.error(f"Failed to drop column '{col_name}' from 'skills'. SQLite might not support drop column: {e}")

    # --- Area 1: Agent Infrastructure Tables ---
        if not inspector.has_table("agent_templates"):
            logger.info("Creating table 'agent_templates'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS agent_templates (
                        id TEXT PRIMARY KEY,
                        name TEXT NOT NULL,
                        description TEXT,
                        system_prompt_path TEXT,
                        max_loop_iterations INTEGER DEFAULT 20
                    )
                """))
                conn.commit()
                logger.info("Table 'agent_templates' created.")
            except Exception as e:
                logger.error(f"Failed to create 'agent_templates': {e}")
        else:
            # Table exists — ensure all columns are present
            template_columns = [c["name"] for c in inspector.get_columns("agent_templates")]
            template_required_columns = [
                ("co_worker_quality_gate", "INTEGER DEFAULT 0"),
                ("rework_threshold", "INTEGER DEFAULT 80"),
                ("max_rework_attempts", "INTEGER DEFAULT 3"),
                ("user_id", "TEXT")
            ]
            for col_name, col_type in template_required_columns:
                if col_name not in template_columns:
                    logger.info(f"Adding column '{col_name}' to 'agent_templates' table...")
                    try:
                        conn.execute(text(f"ALTER TABLE agent_templates ADD COLUMN {col_name} {col_type}"))
                        conn.commit()
                        logger.info(f"Successfully added '{col_name}' to agent_templates.")
                    except Exception as e:
                        logger.error(f"Failed to add column '{col_name}' to agent_templates: {e}")
                
        if not inspector.has_table("agent_instances"):
            logger.info("Creating table 'agent_instances'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS agent_instances (
                        id TEXT PRIMARY KEY,
                        template_id TEXT NOT NULL REFERENCES agent_templates(id),
                        session_id INTEGER REFERENCES sessions(id),
                        mesh_node_id TEXT,
                        status TEXT DEFAULT 'idle',
                        current_workspace_jail TEXT,
                        last_heartbeat DATETIME DEFAULT CURRENT_TIMESTAMP,
                        total_runs INTEGER DEFAULT 0,
                        successful_runs INTEGER DEFAULT 0,
                        total_tokens_accumulated INTEGER DEFAULT 0,
                        total_input_tokens INTEGER DEFAULT 0,
                        total_output_tokens INTEGER DEFAULT 0,
                        total_running_time_seconds INTEGER DEFAULT 0,
                        tool_call_counts TEXT DEFAULT '{}',
                        last_reasoning TEXT,
                        last_error TEXT
                    )
                """))
                conn.commit()
                logger.info("Table 'agent_instances' created.")
            except Exception as e:
                logger.error(f"Failed to create 'agent_instances': {e}")
        else:
            # Multi-iteration: Table exists — ensure all metrics & diagnostic columns are present
            instance_columns = [c["name"] for c in inspector.get_columns("agent_instances")]
            instance_required_columns = [
                ("total_runs", "INTEGER DEFAULT 0"),
                ("successful_runs", "INTEGER DEFAULT 0"),
                ("total_tokens_accumulated", "INTEGER DEFAULT 0"),
                ("total_input_tokens", "INTEGER DEFAULT 0"),
                ("total_output_tokens", "INTEGER DEFAULT 0"),
                ("total_running_time_seconds", "INTEGER DEFAULT 0"),
                ("tool_call_counts", "TEXT DEFAULT '{}'"),
                ("last_reasoning", "TEXT"),
                ("last_error", "TEXT"),
                ("evaluation_status", "TEXT"),
                ("current_rework_attempt", "INTEGER DEFAULT 0"),
                ("latest_quality_score", "INTEGER"),
                ("user_id", "TEXT")
            ]
            for col_name, col_type in instance_required_columns:
                if col_name not in instance_columns:
                    logger.info(f"Adding column '{col_name}' to 'agent_instances' table...")
                    try:
                        conn.execute(text(f"ALTER TABLE agent_instances ADD COLUMN {col_name} {col_type}"))
                        conn.commit()
                        logger.info(f"Successfully added '{col_name}' to agent_instances.")
                    except Exception as e:
                        logger.error(f"Failed to add column '{col_name}' to agent_instances: {e}")
                
        if not inspector.has_table("agent_triggers"):
            logger.info("Creating table 'agent_triggers'...")
            try:
                conn.execute(text("""
                    CREATE TABLE IF NOT EXISTS agent_triggers (
                        id TEXT PRIMARY KEY,
                        instance_id TEXT NOT NULL REFERENCES agent_instances(id),
                        trigger_type TEXT NOT NULL,
                        cron_expression TEXT,
                        interval_seconds INTEGER,
                        default_prompt TEXT,
                        webhook_secret TEXT,
                        webhook_mapping_schema JSON
                    )
                """))
                conn.commit()
                logger.info("Table 'agent_triggers' created.")
            except Exception as e:
                logger.error(f"Failed to create 'agent_triggers': {e}")
        else:
            # Table exists — ensure all columns are present
            trigger_columns = [c["name"] for c in inspector.get_columns("agent_triggers")]
            trigger_required_columns = [
                ("interval_seconds", "INTEGER"),
                ("default_prompt",   "TEXT"),
                ("webhook_secret",   "TEXT"),
                ("webhook_mapping_schema", "JSON")
            ]
            for col_name, col_type in trigger_required_columns:
                if col_name not in trigger_columns:
                    logger.info(f"Adding column '{col_name}' to 'agent_triggers' table...")
                    try:
                        conn.execute(text(f"ALTER TABLE agent_triggers ADD COLUMN {col_name} {col_type}"))
                        conn.commit()
                    except Exception as e:
                        logger.error(f"Failed to add column '{col_name}' to agent_triggers: {e}")

        # --- Ensure user_id is populated for existing agents ---
        try:
            # Find a default user (first one)
            first_user = conn.execute(text("SELECT id FROM users LIMIT 1")).fetchone()
            if first_user:
                default_user_id = first_user[0]
                conn.execute(text(f"UPDATE agent_instances SET user_id = '{default_user_id}' WHERE user_id IS NULL"))
                conn.execute(text(f"UPDATE agent_templates SET user_id = '{default_user_id}' WHERE user_id IS NULL"))
                conn.commit()
                logger.info(f"Assigned existing agents to user '{default_user_id}'.")
        except Exception as e:
            logger.warning(f"Could not assign default user to existing agents: {e}")

    logger.info("Database migrations complete.")




if __name__ == "__main__":
    run_migrations()
