main.py

import os
import json
import logging
import boto3
import openai
from typing import Optional
from src.state import VideoGenerationState
from src.nodes.slide_creation.Slide_Creation_Node_refactor import run
from src.nodes.system_design.system_design_pipeline import SystemDesignPipelineNode
from tracker_utils_tl.tracker import update_session, save_session, save_course

# Initialize boto3 client
s3_client = boto3.client("s3")
BUCKET_NAME = "tech-learn-state"
STATE_FILE_PREFIX = "video_states/"

# LLM CONFIGURATION
openai.api_key = os.environ.get("OPENAI_API_KEY")
LLM_CONFIDENCE_THRESHOLD = 0.6

# Logger Setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# LLM TOPIC ANALYZER
def analyze_topic_with_llm(topic: str) -> Optional[str]:
    """
    LLM analyzes ONLY the topic to decide routing.
    Returns: SYSTEM_DESIGN | SLIDE_CREATION | None
    
    NOTE: This should ONLY be called as a last resort fallback.
    Metadata from UI should always take priority.
    """

    if not openai.api_key:
        logger.warning("OPENAI_API_KEY not set, skipping LLM analysis")
        return None

    prompt = f"""
You are an AI routing agent.

Classify the topic into one of the following:
- SYSTEM_DESIGN → architecture, scalability, components, distributed systems,
                   "Design X", load balancer, database, cache, microservices
- SLIDE_CREATION → programming tutorials, syntax, concepts, algorithms

Respond ONLY in valid JSON.

Topic:
{topic}

Output:
{{
  "intent": "SYSTEM_DESIGN" | "SLIDE_CREATION",
  "confidence": 0.0 to 1.0
}}
"""

    try:
        # Use new OpenAI SDK v1.0+ format
        from openai import OpenAI
        client = OpenAI()  # Uses OPENAI_API_KEY from environment
        
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
        )

        result = json.loads(response.choices[0].message.content)
        intent = result.get("intent")
        confidence = float(result.get("confidence", 0))

        logger.info(f"LLM decision: {intent} (confidence={confidence})")

        if confidence >= LLM_CONFIDENCE_THRESHOLD:
            return intent

    except Exception as e:
        logger.warning(f"LLM topic analysis failed: {e}")

    return None


def detect_design_level_with_ai(state: VideoGenerationState) -> Optional[str]:
    """
    Use AI to determine if topic is HLD (High-Level Design) or LLD (Low-Level Design).
    
    Returns: "HLD" | "LLD" (defaults to "HLD" on failure)
    """
    
    if not openai.api_key:
        logger.warning("OPENAI_API_KEY not set, defaulting to HLD")
        return "HLD"
    
    topic = state.topic or ""
    title = ""
    if state.optional_params:
        title = state.optional_params.get("chapter_title", "")
    
    prompt = f"""
You are an expert system design educator. Analyze the following topic and classify it as EITHER:
- **HLD (High-Level Design)**: System architecture, scalability, distributed systems, APIs, data flow, infrastructure
- **LLD (Low-Level Design)**: Class design, code implementation, design patterns, algorithms, SOLID principles, data structures

IMPORTANT: You MUST choose ONLY ONE - either "HLD" or "LLD". Do NOT return "BOTH".

TOPIC: {topic}
TITLE: {title}

CLASSIFICATION RULES:
1. If topic mentions "Design [System]" (e.g., "Design Twitter", "Design Uber") → HLD
2. If topic mentions a design pattern name (e.g., "Singleton", "Factory", "Observer") → LLD
3. If topic mentions algorithms or data structures (e.g., "LRU Cache", "Trie") → LLD
4. If topic mentions system components or architecture → HLD
5. If topic mentions class diagrams, OOP, or code structure → LLD
6. When in doubt, prefer HLD

EXAMPLES:
- "Design Twitter" → HLD (system architecture)
- "Factory Pattern" → LLD (design pattern implementation)
- "Design URL Shortener" → HLD (distributed system)
- "Singleton Pattern" → LLD (code implementation)
- "Design Uber" → HLD (distributed system)
- "Observer Pattern" → LLD (design pattern)
- "Design Instagram" → HLD (scalable architecture)
- "LRU Cache Implementation" → LLD (algorithm + data structure)
- "Adapter Pattern" → LLD (design pattern)
- "Strategy Pattern" → LLD (design pattern)
- "Load Balancing" → HLD (system component)
- "Database Sharding" → HLD (architecture)

Respond ONLY in valid JSON (choose ONLY "HLD" or "LLD"):
{{
  "design_level": "HLD" | "LLD",
  "confidence": 0.0 to 1.0,
  "reasoning": "Brief explanation"
}}
"""
    
    try:
        from openai import OpenAI
        client = OpenAI()  # Uses OPENAI_API_KEY from environment
        
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": "You are a system design expert. Output ONLY valid JSON without markdown. You MUST choose either HLD or LLD, never BOTH."
                },
                {"role": "user", "content": prompt}
            ],
            temperature=0.2,
        )
        
        result = json.loads(response.choices[0].message.content)
        design_level = result.get("design_level", "HLD")
        confidence = float(result.get("confidence", 0))
        reasoning = result.get("reasoning", "")
        
        # Validate that it's only HLD or LLD
        if design_level not in ["HLD", "LLD"]:
            logger.warning(f"AI returned invalid design_level '{design_level}', defaulting to HLD")
            design_level = "HLD"
        
        logger.info(f"🎯 AI Design Level Detection: {design_level} (confidence={confidence:.2f})")
        logger.info(f"   Reasoning: {reasoning}")
        
        # Only use AI result if confidence is high
        if confidence >= 0.7:
            return design_level
        else:
            logger.warning(f"Low confidence ({confidence:.2f}), defaulting to HLD")
            return "HLD"
            
    except Exception as e:
        logger.error(f"AI design level detection failed: {e}")
        return "HLD"


# AGENT ROUTER
def select_pipeline_node(state: VideoGenerationState) -> str:
    logger.info(f"DEBUG: Entering select_pipeline_node. state.topic: {state.topic}")
    """
    Agent routing logic with METADATA-FIRST approach.
    
    STRONG SYSTEM DESIGN CHECK:
    Routes to SYSTEM_DESIGN pipeline if ANY of these conditions are met:
    - programming_language in ["system_design", "system design"]
    - optional_params.course_type == "system_design"
    - topic_analysis.type == "system_design"
    - course_name contains "system design"
    
    This ensures UI metadata is ALWAYS trusted over LLM classification.
    """

    script_type = (state.optional_params or {}).get("script_type", "")
    if script_type == "consultation":
        logger.info("Routing to CONSULTATION pipeline")
        return "CONSULTATION"

    # Normalize and extract metadata
    course_name = (state.course_name or "Unknown Course").strip().lower()
    language = (state.programming_language or "").strip().lower()
    optional_params = state.optional_params or {}
    
    logger.info(f"🔍 Routing Analysis:")
    logger.info(f"   Course: '{course_name}'")
    logger.info(f"   Language: '{language}'")
    logger.info(f"   Optional Params: {optional_params}")

    # ========== STRONG SYSTEM DESIGN CHECK ==========
    # Extract all system design indicators
    topic_analysis = optional_params.get("topic_analysis", {})
    topic_analysis_type = ""
    if isinstance(topic_analysis, dict):
        topic_analysis_type = (topic_analysis.get("type") or "").strip().lower().replace("_", " ")
    
    course_type = (optional_params.get("course_type") or "").strip().lower().replace("_", " ")
    language_normalized = language.replace("_", " ")
    
    # CONSOLIDATED CHECK: If ANY system design indicator is present -> SYSTEM_DESIGN
    if (
        language in ["system_design", "system design"]
        or language_normalized == "system design"
        or course_type == "system design"
        or topic_analysis_type == "system design"
        or "system design" in course_name
    ):
        # Determine which condition matched
        matched_conditions = []
        if language in ["system_design", "system design"]:
            matched_conditions.append(f"programming_language='{language}'")
        if course_type == "system design":
            matched_conditions.append(f"course_type='{course_type}'")
        if topic_analysis_type == "system design":
            matched_conditions.append(f"topic_analysis.type='{topic_analysis_type}'")
        if "system design" in course_name:
            matched_conditions.append(f"course_name contains 'system design'")
        
        logger.info(f"✅ SYSTEM DESIGN DETECTED: {', '.join(matched_conditions)}")
        
        # Detect design level (HLD vs LLD)
        design_level = detect_design_level_with_ai(state)
        if not state.optional_params:
            state.optional_params = {}
        state.optional_params["design_level"] = design_level
        
        logger.info(f"   Design Level: {design_level}")
        logger.info(f"   Routing to: SYSTEM_DESIGN")
        return "SYSTEM_DESIGN"
    
    # ========== CODING COURSE CHECKS ==========
    if "dsa" in course_name or "language" in course_name:
        logger.info("✅ Routing Rule: course_name implies coding -> SLIDE_CREATION")
        return "SLIDE_CREATION"

    # Check for explicit coding languages
    coding_languages = {
        "python", "java", "javascript", "typescript", "cpp", "c++", "c", 
        "c#", "csharp", "go", "golang", "ruby", "rust", "swift", "kotlin", "php"
    }
    if language in coding_languages:
        logger.info(f"✅ Routing Rule: Known coding language '{language}' -> SLIDE_CREATION")
        return "SLIDE_CREATION"

    # ========== LLM FALLBACK (LAST RESORT) ==========
    logger.warning("⚠️  No metadata match found. Falling back to LLM classification...")
    llm_intent = analyze_topic_with_llm(state.topic)
    if llm_intent:
        logger.info(f"🤖 LLM routing decision: {llm_intent}")
        return llm_intent

    # ========== FINAL FALLBACK ==========
    logger.warning("⚠️  All routing methods failed. Defaulting to SLIDE_CREATION")
    return "SLIDE_CREATION"


def main():
    logger.info("Starting Video Generation Agent")

    # 1. Load input from environment variable
    input_payload = os.environ.get("INPUT_PAYLOAD")
    if not input_payload:
        logger.error("Missing INPUT_PAYLOAD")
        exit(1)

    try:
        event = json.loads(input_payload)
    except json.JSONDecodeError as e:
        logger.error(f"Invalid JSON in INPUT_PAYLOAD: {e}")
        exit(1)

    # 2. Extract session_id
    session_id = extract_session_id(event)
    if not session_id:
        logger.error("session_id not found")
        exit(1)

    # 3. Load state (from event body first, then S3 fallback)
    session_id = extract_session_id(event)
    if not session_id:
        logger.error("session_id not found in input.")
        exit(1)

    # 4. Load state from event body if available
    state_data = None
    if "body" in event:
        try:
            body_data = json.loads(event["body"]) if isinstance(event["body"], str) else event["body"]
            if "topic" in body_data and "session_id" in body_data:
                state_data = body_data
                logger.info("Loaded state data from event body")
        except Exception:
            pass

    # 5. Build S3 key and fallback to S3 if body state is missing
    state_s3_key = f"{STATE_FILE_PREFIX}{session_id}/{session_id}.json"
    if not state_data:
        logger.info(f"Loading state from S3: {state_s3_key}")
        try:
            state_data = load_state_from_s3(state_s3_key)
        except Exception as e:
            logger.error(f"Failed to load state from S3: {e}")
            exit(1)

     # 6. Extract optional course and topic metadata
    optional_params = state_data.get("optional_params") or {}
    course_id = optional_params.get("course_id", 0)
    topic_id = optional_params.get("topic_id", 0)
    chapter_id = optional_params.get("chapter_id", 0)
    topic_title = optional_params.get("chapter_title", "")
    course_name = optional_params.get("course_name", f"Course {course_id}")
    
    # Ensure course_name always has a valid string value (never null/empty)
    if not course_name or course_name.strip() == "":
        course_name = f"Course {course_id}"
    
    # Inject course_name into state data if missing
    if "course_name" not in state_data or not state_data.get("course_name"):
        state_data["course_name"] = course_name

# 7. Persist course information for tracking and analytics
    try:
        save_course(
            course_detail={
                "course_id": int(course_id),
                "topic_id": int(topic_id),
                "chapter_id": int(chapter_id),
                "total_videos": int(optional_params.get("total_videos", 0)),
            },
            course_name=course_name,
        )
    except Exception as e:
        logger.warning(f"Course save skipped: {e}")

    # 8. Create strongly-typed state object
    state = VideoGenerationState(**state_data)

    logger.info(f"Topic: {state.topic}")
    logger.info(f"Programming Language: {state.programming_language}")

    # 9. Let agent decide which pipeline node to run
    selected_node = select_pipeline_node(state)
    logger.info(f"Agent selected node: {selected_node}")

    # 10. Execute the selected pipeline node
    try:
        if selected_node == "SYSTEM_DESIGN":
            node_name = "SystemDesign"

            # Mark system design node as started
            update_session(session_id, node=node_name, status="STARTED")

            # Run system design pipeline
            node = SystemDesignPipelineNode()
            updated_state = node.generate_system_design_video(state)

        elif selected_node == "SLIDE_CREATION":
            node_name = "SlideCreation"

            # Mark slide creation node as started
            update_session(session_id, node=node_name, status="STARTED")

            # Run slide creation pipeline
            updated_state = run(state)
        elif selected_node == "CONSULTATION":
            node_name = "Consultation"
            update_session(session_id, node=node_name, status="STARTED")
            from src.nodes.consultation.consultation_pipeline import ConsultationPipelineNode
            node = ConsultationPipelineNode()
            updated_state = node.generate_consultation_video(state)

        else:
            raise ValueError(f"Unknown node selected: {selected_node}")

        # 11. Save updated execution state back to S3
        save_state_to_s3(state_s3_key, updated_state.model_dump())

        # 12. Mark pipeline execution as completed
        update_session(session_id, node=node_name, status="COMPLETED")
        save_session(
            session_id=session_id,
            course_id=course_id,
            topic_id=topic_id,
            topic_title=topic_title,
            node=node_name,
            status="COMPLETED",
        )

        # 13. Exit successfully if terminal status is reached
        if updated_state.status in [
            "slide_video_generated",
            "system_design_video_generated",
        ]:
            logger.info(f"{node_name} pipeline completed successfully")
            exit(0)

        # 14. Exit with failure if unexpected status is returned
        logger.warning(f"{node_name} finished with status: {updated_state.status}")
        exit(1)

    except Exception as e:
        # 15. Handle pipeline failure and update tracker
        logger.error(f"Pipeline failed: {e}")

        try:
            update_session(session_id, node=node_name, status="FAILED")
            save_session(
                session_id=session_id,
                course_id=course_id,
                topic_id=topic_id,
                topic_title=topic_title,
                node=node_name,
                status="FAILED",
            )
        except Exception:
            pass

        exit(1)

def extract_session_id(event):
    if "session_id" in event:
        return event["session_id"]
    if "headers" in event and "session_id" in event["headers"]:
        return event["headers"]["session_id"]
    if "body" in event:
        try:
            body = json.loads(event["body"]) if isinstance(event["body"], str) else event["body"]
            if "session_id" in body:
                return body["session_id"]
        except Exception:
            pass
    return None

def load_state_from_s3(key: str):
    obj = s3_client.get_object(Bucket=BUCKET_NAME, Key=key)
    return json.loads(obj["Body"].read().decode("utf-8"))


def save_state_to_s3(key: str, data: dict):
    s3_client.put_object(
        Bucket=BUCKET_NAME,
        Key=key,
        Body=json.dumps(data, indent=2),
        ContentType="application/json",
    )

if __name__ == "__main__":
    main()