#!/usr/bin/env python3
"""
main.py
AIVA Command Bridge API - Telnyx Webhook Handlers
"""

import os
import json
import logging
from contextlib import contextmanager
from datetime import datetime
from typing import Optional, Dict, Any, List

import psycopg2
from psycopg2 import pool, extras
from fastapi import FastAPI, HTTPException, Depends, Header, Request, status
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import uvicorn

# Configuration
DB_CONFIG = {
    "host": "postgresql-genesis-u50607.vm.elestio.app",
    "port": 25432,
    "user": "postgres",
    "password": "CiBjh6LM7Yuqkq-jo2r7eQDw",
    "database": "postgres"
}
SCHEMA = "genesis_bridge"
API_KEY = os.getenv("BRIDGE_API_KEY", "genesis-aiva-bridge-key-2024")
TELNYX_WEBHOOK_SECRET = os.getenv("TELNYX_WEBHOOK_SECRET", "")

# Logging Configuration
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("aiva_bridge.log")
    ]
)
logger = logging.getLogger(__name__)

# Database Connection Pool
db_pool: Optional[pool.ThreadedConnectionPool] = None

def initialize_db_pool():
    """Initialize PostgreSQL connection pool"""
    global db_pool
    try:
        db_pool = psycopg2.pool.ThreadedConnectionPool(
            minconn=2,
            maxconn=20,
            **DB_CONFIG
        )
        logger.info("Database connection pool initialized")
    except Exception as e:
        logger.critical(f"Failed to initialize database pool: {e}")
        raise

@contextmanager
def get_db_connection():
    """Context manager for database connections"""
    if not db_pool:
        raise HTTPException(status_code=503, detail="Database not available")
    
    conn = None
    try:
        conn = db_pool.getconn()
        yield conn
    except psycopg2.Error as e:
        logger.error(f"Database error: {e}")
        if conn:
            conn.rollback()
        raise HTTPException(status_code=500, detail="Database error")
    finally:
        if conn:
            db_pool.putconn(conn)

def init_database_schema():
    """Initialize database schema and tables"""
    schema_sql = f"""
    CREATE SCHEMA IF NOT EXISTS {SCHEMA};
    
    CREATE TABLE IF NOT EXISTS {SCHEMA}.directives (
        id SERIAL PRIMARY KEY,
        directive_text TEXT NOT NULL,
        priority INTEGER DEFAULT 5 CHECK (priority BETWEEN 1 AND 10),
        directive_type VARCHAR(50) CHECK (directive_type IN ('task', 'query', 'status_request', 'urgent')),
        status VARCHAR(50) DEFAULT 'pending',
        created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
        processed_at TIMESTAMP WITH TIME ZONE,
        telnyx_conversation_id VARCHAR(255),
        tool_call_id VARCHAR(255),
        metadata JSONB DEFAULT '{{}}'::jsonb
    );
    
    CREATE TABLE IF NOT EXISTS {SCHEMA}.claude_status (
        id INTEGER PRIMARY KEY DEFAULT 1 CHECK (id = 1),
        current_task TEXT,
        task_status VARCHAR(50) DEFAULT 'idle',
        last_updated TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
        task_start_time TIMESTAMP WITH TIME ZONE,
        details JSONB DEFAULT '{{}}'::jsonb
    );
    
    INSERT INTO {SCHEMA}.claude_status (id, current_task, task_status) 
    VALUES (1, 'System initialized', 'idle')
    ON CONFLICT (id) DO NOTHING;
    
    CREATE INDEX IF NOT EXISTS idx_directives_status_created 
        ON {SCHEMA}.directives(status, created_at DESC);
    CREATE INDEX IF NOT EXISTS idx_directives_conversation 
        ON {SCHEMA}.directives(telnyx_conversation_id);
    """
    
    try:
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                cur.execute(schema_sql)
                conn.commit()
                logger.info("Database schema initialized successfully")
    except Exception as e:
        logger.error(f"Schema initialization failed: {e}")
        raise

# Authentication
async def verify_api_key(x_api_key: Optional[str] = Header(None, alias="X-API-Key")):
    """Verify API Key header"""
    if not x_api_key or x_api_key != API_KEY:
        logger.warning(f"Invalid API key attempt: {x_api_key[:10]}..." if x_api_key else "No API key")
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing X-API-Key"
        )
    return x_api_key

# Pydantic Models
class TelnyxToolCall(BaseModel):
    id: str
    name: str
    arguments: str  # JSON string from Telnyx

class TelnyxWebhookPayload(BaseModel):
    conversation_id: Optional[str] = None
    tool_calls: List[TelnyxToolCall]

class DirectiveResponse(BaseModel):
    status: str
    directive_id: Optional[int] = None
    message: str

# FastAPI Application
app = FastAPI(
    title="AIVA Command Bridge",
    description="Telnyx Voice Assistant Bridge for Genesis",
    version="1.0.0",
    docs_url="/docs" if os.getenv("ENV") != "production" else None
)

@app.on_event("startup")
async def startup_event():
    """Initialize application on startup"""
    try:
        initialize_db_pool()
        init_database_schema()
        logger.info("AIVA Bridge started successfully")
    except Exception as e:
        logger.critical(f"Startup failed: {e}")
        raise

@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown"""
    global db_pool
    if db_pool:
        db_pool.closeall()
        logger.info("Database connections closed")

@app.post("/bridge/telnyx/relay_directive", response_model=Dict[str, Any])
async def relay_directive_to_claude(
    request: Request,
    api_key: str = Depends(verify_api_key)
):
    """
    Webhook endpoint for Telnyx relay_directive tool.
    Receives Kinan's directive and queues it for Claude Code.
    """
    try:
        body = await request.json()
        logger.info(f"Received relay_directive webhook: {json.dumps(body)}")
        
        # Parse payload (handle both wrapped and unwrapped formats)
        data = body.get("data", body)
        tool_calls = data.get("tool_calls", [])
        
        if not tool_calls:
            logger.error("No tool calls in webhook payload")
            raise HTTPException(status_code=400, detail="No tool calls found")
        
        tool_call = tool_calls[0]
        tool_call_id = tool_call.get("id")
        conversation_id = data.get("conversation_id", "unknown")
        
        # Parse arguments (Telnyx sends as JSON string)
        try:
            args_str = tool_call.get("arguments", "{}")
            if isinstance(args_str, str):
                arguments = json.loads(args_str)
            else:
                arguments = args_str
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse arguments: {e}")
            raise HTTPException(status_code=400, detail="Invalid arguments format")
        
        directive_text = arguments.get("directive_text")
        if not directive_text:
            raise HTTPException(status_code=400, detail="directive_text is required")
        
        priority = arguments.get("priority", 5)
        directive_type = arguments.get("directive_type", "task")
        
        # Insert into database
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                cur.execute(f"""
                    INSERT INTO {SCHEMA}.directives 
                    (directive_text, priority, directive_type, status, 
                     telnyx_conversation_id, tool_call_id, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s, CURRENT_TIMESTAMP)
                    RETURNING id
                """, (
                    directive_text, 
                    priority, 
                    directive_type, 
                    "pending",
                    conversation_id,
                    tool_call_id
                ))
                
                directive_id = cur.fetchone()[0]
                conn.commit()
                
                logger.info(f"Directive {directive_id} queued: {directive_text[:50]}...")
        
        # Return Telnyx-compatible response
        return {
            "data": {
                "tool_call_results": [
                    {
                        "id": tool_call_id,
                        "output": f"Directive received (ID: {directive_id}). Claude Code will process this {directive_type} with priority {priority}."
                    }
                ]
            }
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error in relay_directive: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail="Internal server error")

@app.post("/bridge/telnyx/check_status", response_model=Dict[str, Any])
async def check_claude_status(
    request: Request,
    api_key: str = Depends(verify_api_key)
):
    """
    Webhook endpoint for Telnyx check_claude_status tool.
    Returns current Claude Code activity status.
    """
    try:
        body = await request.json()
        logger.info(f"Received check_status webhook: {json.dumps(body)}")
        
        data = body.get("data", body)
        tool_calls = data.get("tool_calls", [])
        
        if not tool_calls:
            raise HTTPException(status_code=400, detail="No tool calls found")
        
        tool_call = tool_calls[0]
        tool_call_id = tool_call.get("id")
        
        # Parse arguments
        try:
            args_str = tool_call.get("arguments", "{}")
            if isinstance(args_str, str):
                arguments = json.loads(args_str) if args_str else {}
            else:
                arguments = args_str
        except json.JSONDecodeError:
            arguments = {}
        
        include_history = arguments.get("include_history", False)
        history_limit = min(arguments.get("history_limit", 5), 50)
        
        with get_db_connection() as conn:
            with conn.cursor(cursor_factory=extras.RealDictCursor) as cur:
                # Get current status
                cur.execute(f"""
                    SELECT current_task, task_status, last_updated, task_start_time 
                    FROM {SCHEMA}.claude_status 
                    WHERE id = 1
                """)
                status_row = cur.fetchone()
                
                if not status_row:
                    current_task = "Unknown"
                    task_status = "unknown"
                    last_updated = datetime.now()
                else:
                    current_task = status_row["current_task"] or "Idle"
                    task_status = status_row["task_status"] or "idle"
                    last_updated = status_row["last_updated"]
                
                # Calculate duration if working
                duration_str = ""
                if task_status == "working" and status_row and status_row["task_start_time"]:
                    duration = datetime.now() - status_row["task_start_time"]
                    duration_str = f" for {duration.seconds // 60} minutes"
                
                result_text = f"Claude Code is currently {task_status}{duration_str}."
                
                if current_task and current_task != "idle":
                    result_text += f" Current focus: {current_task}."
                
                # Add history if requested
                if include_history:
                    cur.execute(f"""
                        SELECT directive_text, directive_type, status, 
                               created_at, priority
                        FROM {SCHEMA}.directives 
                        ORDER BY created_at DESC 
                        LIMIT %s
                    """, (history_limit,))
                    
                    history = cur.fetchall()
                    if history:
                        result_text += f" Recent activity ({len(history)} items):"
                        for i, item in enumerate(history, 1):
                            status_icon = "✓" if item["status"] == "completed" else "○"
                            text_preview = item["directive_text"][:40] + "..." if len(item["directive_text"]) > 40 else item["directive_text"]
                            result_text += f" [{i}] {status_icon} {item['directive_type']} (P{item['priority']}): {text_preview}."
        
        return {
            "data": {
                "tool_call_results": [
                    {
                        "id": tool_call_id,
                        "output": result_text
                    }
                ]
            }
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in check_status: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail="Internal server error")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    try:
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                cur.execute("SELECT 1")
        return {"status": "healthy", "service": "aiva-bridge", "database": "connected"}
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        raise HTTPException(status_code=503, detail="Service unavailable")

@app.get("/bridge/directives/pending")
async def get_pending_directives(
    limit: int = 10,
    api_key: str = Depends(verify_api_key)
):
    """
    Endpoint for Claude Code to poll pending directives.
    Returns list of pending directives ordered by priority and time.
    """
    try:
        with get_db_connection() as conn:
            with conn.cursor(cursor_factory=extras.RealDictCursor) as cur:
                cur.execute(f"""
                    SELECT id, directive_text, priority, directive_type, 
                           created_at, telnyx_conversation_id
                    FROM {SCHEMA}.directives 
                    WHERE status = 'pending'
                    ORDER BY priority DESC, created_at ASC
                    LIMIT %s
                """, (limit,))
                
                directives = cur.fetchall()
                
                # Mark as retrieved (optional state)
                if directives:
                    ids = [d["id"] for d in directives]
                    cur.execute(f"""
                        UPDATE {SCHEMA}.directives 
                        SET status = 'retrieved', 
                            metadata = jsonb_set(metadata, '{{retrieved_at}}', to_jsonb(CURRENT_TIMESTAMP))
                        WHERE id = ANY(%s)
                    """, (ids,))
                    conn.commit()
                
                return {
                    "directives": [dict(d) for d in directives],
                    "count": len(directives)
                }
                
    except Exception as e:
        logger.error(f"Error fetching directives: {e}")
        raise HTTPException(status_code=500, detail="Failed to fetch directives")

@app.post("/bridge/status/update")
async def update_claude_status(
    status_update: Dict[str, Any],
    api_key: str = Depends(verify_api_key)
):
    """
    Endpoint for Claude Code to update its current status.
    """
    try:
        current_task = status_update.get("current_task", "idle")
        task_status = status_update.get("task_status", "idle")
        details = status_update.get("details", {})
        
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                # Use task_start_time tracking
                time_sql = ""
                if task_status == "working":
                    time_sql = ", task_start_time = COALESCE(task_start_time, CURRENT_TIMESTAMP)"
                elif task_status == "idle":
                    time_sql = ", task_start_time = NULL"
                
                cur.execute(f"""
                    INSERT INTO {SCHEMA}.claude_status (id, current_task, task_status, last_updated, details)
                    VALUES (1, %s, %s, CURRENT_TIMESTAMP, %s)
                    ON CONFLICT (id) 
                    DO UPDATE SET 
                        current_task = EXCLUDED.current_task,
                        task_status = EXCLUDED.task_status,
                        last_updated = CURRENT_TIMESTAMP,
                        details = EXCLUDED.details
                        {time_sql}
                """, (current_task, task_status, json.dumps(details)))
                
                conn.commit()
        
        return {"status": "updated", "current_task": current_task, "task_status": task_status}
        
    except Exception as e:
        logger.error(f"Error updating status: {e}")
        raise HTTPException(status_code=500, detail="Failed to update status")

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        log_level="info",
        workers=4
    )