from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import pandas as pd
import psycopg2
from psycopg2.extras import RealDictCursor
import os
from io import BytesIO

app = FastAPI()

# Database connection parameters
DB_CONFIG = {
    "host": "13.126.242.31",
    "database": "aml",
    "user": "dev_cbs_admin",
    "password": "Finovate@2023"
}

def get_db_connection():
    """Create database connection"""
    return psycopg2.connect(**DB_CONFIG, cursor_factory=RealDictCursor)

def read_columns_from_file(template_id: str):
    """Read column names from output file"""
    filename = f'output_{template_id}.txt'
    try:
        with open(filename, 'r') as f:
            content = f.read()
            # Safely evaluate the string representation of the list
            columns = eval(content)
        return columns
    except Exception as e:
        raise HTTPException(status_code=404, detail=f"Error reading columns from file: {str(e)}")

@app.get("/download/{template_id}")
async def download_excel(template_id: str):
    """
    Download data as Excel file based on template ID
    """
    # Read columns from selected file
    columns = read_columns_from_file(template_id)
    
    if not columns:
        raise HTTPException(status_code=404, detail="No columns found in output file")
    
    try:
        # Connect to database
        conn = get_db_connection()
        cur = conn.cursor()
        
        # Get existing columns from the table
        cur.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'large_transactions'
        """)
        existing_columns = {row['column_name'].lower(): row['column_name'] for row in cur.fetchall()}
        
        # Filter out non-existent columns and use actual column names from DB
        valid_columns = []
        for col in columns:
            if col.lower() in existing_columns:
                valid_columns.append(existing_columns[col.lower()])
        
        if not valid_columns:
            raise HTTPException(status_code=404, detail="None of the specified columns exist in the table")
            
        # Create SQL query with proper column quoting
        quoted_columns = [f'"{col}"' for col in valid_columns]
        query = f"""
            SELECT {', '.join(quoted_columns)}
            FROM large_transactions
        """
        
        # Execute query
        cur.execute(query)
        rows = cur.fetchall()
        
        # Convert to DataFrame
        df = pd.DataFrame(rows)
        
        # Create Excel file in memory
        output = BytesIO()
        with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
            df.to_excel(writer, index=False, sheet_name='Data')
        output.seek(0)
        
        # Close database connection
        cur.close()
        conn.close()
        
        # Return the Excel file as a downloadable response
        return StreamingResponse(
            output,
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            headers={
                "Content-Disposition": f"attachment; filename=large_transactions_{template_id}.xlsx"
            }
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error fetching data: {str(e)}")

@app.get("/templates")
async def list_templates():
    """List all available template IDs"""
    try:
        output_files = [f.replace('output_', '').replace('.txt', '') 
                       for f in os.listdir('.') 
                       if f.startswith('output_') and f.endswith('.txt')]
        return {"templates": output_files}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error listing templates: {str(e)}")