#!/usr/bin/env python3
"""Fix missing sentiment columns in citation_records table."""

import asyncio
import asyncpg


async def fix_schema():
    # Read DATABASE_URL from .env file in project root
    db_url = "postgresql://chiguyong@localhost:5432/geo_platform"
    
    conn = await asyncpg.connect(db_url)
    try:
        # Check existing columns
        rows = await conn.fetch(
            "SELECT column_name FROM information_schema.columns WHERE table_name = 'citation_records'"
        )
        existing = {r["column_name"] for r in rows}
        print(f"Existing columns: {existing}")
        
        needed = ["sentiment", "sentiment_confidence", "sentiment_key_phrases"]
        missing = [c for c in needed if c not in existing]
        
        if not missing:
            print("All sentiment columns already exist.")
            return
        
        print(f"Missing columns: {missing}")
        
        if "sentiment" in missing:
            await conn.execute(
                "ALTER TABLE citation_records ADD COLUMN sentiment VARCHAR(20) NULL"
            )
            print("Added sentiment column")
        
        if "sentiment_confidence" in missing:
            await conn.execute(
                "ALTER TABLE citation_records ADD COLUMN sentiment_confidence DOUBLE PRECISION NULL"
            )
            print("Added sentiment_confidence column")
        
        if "sentiment_key_phrases" in missing:
            await conn.execute(
                "ALTER TABLE citation_records ADD COLUMN sentiment_key_phrases JSONB NULL"
            )
            print("Added sentiment_key_phrases column")
        
        print("Schema fix complete.")
    finally:
        await conn.close()


if __name__ == "__main__":
    asyncio.run(fix_schema())