import os
import psycopg
from dotenv import load_dotenv

load_dotenv("d:/Mobcoder Pharam Care/.env")
db_url = os.getenv("AKS_DATABASE_URL")
if db_url:
    db_url = db_url.replace("postgresql+psycopg://", "postgresql://")

print(f"Connecting to Postgres...")

try:
    with psycopg.connect(db_url) as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';")
            tables = [row[0] for row in cur.fetchall()]
            
            if 'gateway_answers' not in tables:
                print("Creating 'gateway_answers' table...")
                cur.execute("""
                    CREATE TABLE gateway_answers (
                        answer_id VARCHAR(255) PRIMARY KEY,
                        request_id VARCHAR(255),
                        question TEXT,
                        user_role VARCHAR(255),
                        audience VARCHAR(255),
                        geography VARCHAR(255),
                        therapy_area VARCHAR(255),
                        policy_profile VARCHAR(255),
                        decision VARCHAR(255),
                        policy_outcome VARCHAR(255),
                        retrieval_confidence FLOAT,
                        citation_validation_passed BOOLEAN,
                        embedding_model VARCHAR(255),
                        generation_model VARCHAR(255),
                        response_json JSONB,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    );
                """)
                conn.commit()
                print("Table created successfully!")
            
            # Now fetch count
            cur.execute("SELECT COUNT(*) FROM gateway_answers;")
            count = cur.fetchone()[0]
            print(f"Total chat logs in gateway_answers: {count}")

            if 'audit_logs' in tables:
                cur.execute("SELECT COUNT(*) FROM audit_logs;")
                acount = cur.fetchone()[0]
                print(f"Total audit logs in audit_logs: {acount}")

except Exception as e:
    print(f"Error connecting to Postgres: {e}")