import os
import json
import logging
import random
import re
from langchain_core.messages import SystemMessage, HumanMessage  #
from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFDirectoryLoader, WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_tavily import TavilySearch
from langchain.tools.retriever import create_retriever_tool
from langchain.agents import create_tool_calling_agent, AgentExecutor
import gradio as gr
import numpy as np
from dataclasses import dataclass
from typing import List, Dict
import uuid
import socket
from pythia import scenario_simulator

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Load environment variables
os.environ['USER_AGENT'] = os.getenv('USER_AGENT', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)')

# load_dotenv()
# Verify API keys
openai_api_key = os.getenv("OPENAI_API_KEY")
tavily_api_key = os.getenv("TAVILY_API_KEY")
if not openai_api_key or not tavily_api_key:
    logger.error("Missing API keys: OPENAI_API_KEY or TAVILY_API_KEY")
    raise ValueError("API keys not set in environment")

# Persona-related definitions
@dataclass
class Persona:
    id: str
    name: str
    age: int
    gender: str
    occupation: str
    adoption_stage: str
    pain_points: List[str]
    motivations: List[str]
    demographics: Dict

# Store active persona and conversation state
active_persona = None
conversation_count = 0
MAX_QUESTIONS = 20



def load_web_resources(file='websites.txt'):
    urls = ['https://steveblank.com/']
    if os.path.exists(file):
        with open(file, 'r') as f:
            urls.extend(line.strip() for line in f if line.strip())
    web_docs = []
    for url in urls:
        try:
            loader = WebBaseLoader(url)
            web_docs.extend(loader.load())
            logger.info(f"Loaded web resources from {url}")
        except Exception as e:
            logger.error(f"Failed to load {url}: {e}")
    return web_docs
def persona_generator(user_input: str) -> List[Dict]:
    """Generates five UX personas based on user input and adoption stages."""
    try:
        llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
        
        # Always pull Dr. Sue's material for context
        retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
        docs = retriever.invoke(user_input)
        dr_sue_context = "\n\n".join([doc.page_content for doc in docs])
        logger.info(f"Incorporated Dr. Sue context: {dr_sue_context[:200]}...")

        # Build messages directly
        system_content = f"""
You are a UX persona generator for Dr. Sue Cohen's Commercializing New Technologies class. Parse the user input (which may be narrative, e.g., 'Single mother, tired') into a JSON list of exactly 5 personas, one for each adoption stage: Innovator, Early Adopter, Early Majority, Late Majority, Laggard. Output ONLY valid JSON matching the template below—do NOT add extra text or prose.

Incorporate all input details (e.g., occupation, age, gender, family status, pain points). Convert narrative phrases into concise pain_points and motivations (e.g., 'She is tired' → 'Shift-induced exhaustion'). Use input pain points if provided (e.g., 'perpetually tired' → 'Shift-induced exhaustion'); otherwise, generate stage-specific ones consistent with the input context. Generate unique motivations and scenarios per persona, reflecting their adoption stage and healthcare technology adoption (e.g., AI tools, telehealth). Scenarios show tech use, tied to stage, pain points, and motivations. Include demographics (location, family_status, workplace). Generate unique UUID for 'id' and unique realistic names. Use this course context for depth (e.g., regulatory barriers, scaling): {dr_sue_context}

JSON template:
[
  {{
    "id": "UUID",
    "name": "Unique Name",
    "age": integer,
    "gender": "Gender",
    "occupation": "Occupation",
    "adoption_stage": "Stage",
    "pain_points": ["Phrase1", "Phrase2", "Phrase3"],
    "motivations": ["Phrase1", "Phrase2", "Phrase3"],
    "demographics": {{"location": "Place", "family_status": "Status", "workplace": "Place"}},
    "scenarios": ["Narrative1", "Narrative2"]
  }}
]

Adoption stage traits:
  - Innovators: Risk-takers, tech enthusiasts, embrace new tools.
  - Early Adopters: Visionaries, influence others, accept moderate risk.
  - Early Majority: Pragmatic, seek proven value, efficiency-focused.
  - Late Majority: Skeptical, need social proof, risk-averse.
  - Laggards: Resistant, distrust change, adopt last.
        """.strip()
        
        messages = [
            SystemMessage(content=system_content),
            HumanMessage(content=user_input)
        ]
        
        logger.info(f"Invoking LLM with input: {user_input}")
        response = llm.invoke(messages)
        logger.info(f"LLM response: {response.content[:200]}...")
        personas = json.loads(response.content)
        
        # Validation
        if not isinstance(personas, list) or len(personas) != 5:
            raise ValueError(f"Expected 5 personas, got {len(personas)}")
        adoption_stages = ["Innovator", "Early Adopter", "Early Majority", "Late Majority", "Laggard"]
        for i, persona in enumerate(personas):
            if "id" not in persona or not persona["id"]:
                persona["id"] = str(uuid.uuid4())
            if "adoption_stage" not in persona or persona["adoption_stage"] not in adoption_stages:
                persona["adoption_stage"] = adoption_stages[i]
            required_fields = ["name", "age", "gender", "occupation", "pain_points", "motivations", "demographics", "scenarios"]
            for field in required_fields:
                if field not in persona:
                    raise ValueError(f"Missing field {field} in persona {i}")
            if not isinstance(persona["age"], int):
                try:
                    persona["age"] = int(persona["age"])
                except (ValueError, TypeError):
                    raise ValueError("Invalid age")
            if not isinstance(persona["pain_points"], list) or not all(isinstance(p, str) for p in persona["pain_points"]):
                raise ValueError(f"Invalid pain_points in persona {i}")
            if not isinstance(persona["motivations"], list) or not all(isinstance(p, str) for p in persona["motivations"]):
                raise ValueError(f"Invalid motivations in persona {i}")
            if not isinstance(persona["scenarios"], list) or not all(isinstance(s, str) for s in persona["scenarios"]):
                raise ValueError(f"Invalid scenarios in persona {i}")
            if not isinstance(persona["demographics"], dict) or not all(k in persona["demographics"] for k in ["location", "family_status", "workplace"]):
                raise ValueError(f"Invalid demographics in persona {i}")
        logger.info(f"Successfully generated {len(personas)} personas")
        return personas
    except Exception as e:
        logger.error(f"Persona generation failed: {str(e)}")
        raise ValueError(f"Failed to generate personas: {str(e)}. Use format: 'Persona Mode: Generate Persona: [description]'")
def build_vectorstore(pdf_dir="ConvertedFiles"):
    try:
        loader = PyPDFDirectoryLoader(pdf_dir)
        docs = loader.load()
        logger.info(f"Loaded {len(docs)} documents from {pdf_dir}")
        web_docs = load_web_resources()
        all_docs = docs + web_docs
        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = splitter.split_documents(all_docs)
        logger.info(f"Split into {len(splits)} document chunks")
        embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        vectorstore = FAISS.from_documents(splits, embeddings)
        vectorstore.save_local("faiss_index")
        logger.info("Saved FAISS index to faiss_index/")
        return vectorstore
    except Exception as e:
        logger.error(f"Failed to build vectorstore: {e}")
        raise


    """Generates five UX personas based on user input and adoption stages."""
    try:
        llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
        
        # Always pull Dr. Sue's material for context
        retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
        docs = retriever.invoke(user_input)
        dr_sue_context = "\n\n".join([doc.page_content for doc in docs])
        logger.info(f"Incorporated Dr. Sue context: {dr_sue_context[:200]}...")

        # Build messages directly
        system_content = f"""
You are a UX persona generator for Dr. Sue Cohen's Commercializing New Technologies class. Parse the user input (which may be narrative, e.g., 'Single mother, tired') into a JSON list of exactly 5 personas, one for each adoption stage: Innovator, Early Adopter, Early Majority, Late Majority, Laggard. Output ONLY valid JSON matching the template below—do NOT add extra text or prose.

Incorporate all input details (e.g., occupation, age, gender, family status, pain points). Convert narrative phrases into concise pain_points and motivations (e.g., 'She is tired' → 'Shift-induced exhaustion'). Use input pain points if provided; otherwise, generate stage-specific ones consistent with the input context. Generate unique motivations and scenarios per persona, reflecting their adoption stage and healthcare technology adoption (e.g., AI tools, telehealth). Scenarios show tech use, tied to stage, pain points, and motivations. Include demographics (location, family_status, workplace). Generate unique UUID for 'id' and unique realistic names. Use this course context for depth (e.g., regulatory barriers, scaling): {dr_sue_context}

JSON template:
[
  {{
    "id": "UUID",
    "name": "Unique Name",
    "age": integer,
    "gender": "Gender",
    "occupation": "Occupation",
    "adoption_stage": "Stage",
    "pain_points": ["Phrase1", "Phrase2", "Phrase3"],
    "motivations": ["Phrase1", "Phrase2", "Phrase3"],
    "demographics": {{"location": "Place", "family_status": "Status", "workplace": "Place"}},
    "scenarios": ["Narrative1", "Narrative2"]
  }}
]

Adoption stage traits:
  - Innovators: Risk-takers, tech enthusiasts, embrace new tools.
  - Early Adopters: Visionaries, influence others, accept moderate risk.
  - Early Majority: Pragmatic, seek proven value, efficiency-focused.
  - Late Majority: Skeptical, need social proof, risk-averse.
  - Laggards: Resistant, distrust change, adopt last.
        """.strip()
        
        messages = [
            SystemMessage(content=system_content),
            HumanMessage(content=user_input)
        ]
        
        logger.info(f"Invoking LLM with input: {user_input}")
        response = llm.invoke(messages)
        logger.info(f"LLM response: {response.content[:200]}...")
        personas = json.loads(response.content)
        
        # Validation
        if not isinstance(personas, list) or len(personas) != 5:
            raise ValueError(f"Expected 5 personas, got {len(personas)}")
        adoption_stages = ["Innovator", "Early Adopter", "Early Majority", "Late Majority", "Laggard"]
        for i, persona in enumerate(personas):
            if "id" not in persona or not persona["id"]:
                persona["id"] = str(uuid.uuid4())
            if "adoption_stage" not in persona or persona["adoption_stage"] not in adoption_stages:
                persona["adoption_stage"] = adoption_stages[i]
            required_fields = ["name", "age", "gender", "occupation", "pain_points", "motivations", "demographics", "scenarios"]
            for field in required_fields:
                if field not in persona:
                    raise ValueError(f"Missing field {field} in persona {i}")
            if not isinstance(persona["age"], int):
                try:
                    persona["age"] = int(persona["age"])
                except (ValueError, TypeError):
                    raise ValueError("Invalid age")
            if not isinstance(persona["pain_points"], list) or not all(isinstance(p, str) for p in persona["pain_points"]):
                raise ValueError(f"Invalid pain_points in persona {i}")
            if not isinstance(persona["motivations"], list) or not all(isinstance(p, str) for p in persona["motivations"]):
                raise ValueError(f"Invalid motivations in persona {i}")
            if not isinstance(persona["scenarios"], list) or not all(isinstance(s, str) for s in persona["scenarios"]):
                raise ValueError(f"Invalid scenarios in persona {i}")
            if not isinstance(persona["demographics"], dict) or not all(k in persona["demographics"] for k in ["location", "family_status", "workplace"]):
                raise ValueError(f"Invalid demographics in persona {i}")
        logger.info(f"Successfully generated {len(personas)} personas")
        return personas
    except Exception as e:
        logger.error(f"Persona generation failed: {str(e)}")
        raise ValueError(f"Failed to generate personas: {str(e)}. Use format: 'Persona Mode: Generate Persona: [description]'")
def validate_persona_json(persona_json: str) -> Dict:
    """Validates and parses a persona JSON string."""
    try:
        # Clean JSON input to handle whitespace and line breaks
        persona_json = ' '.join(persona_json.strip().split())
        persona = json.loads(persona_json)
        required_fields = ["id", "name", "age", "gender", "occupation", "adoption_stage", "pain_points", "motivations", "demographics", "scenarios"]
        for field in required_fields:
            if field not in persona:
                raise ValueError(f"Missing field: {field}")
        if not isinstance(persona["age"], int):
            try:
                persona["age"] = int(persona["age"])
            except (ValueError, TypeError):
                raise ValueError("age must be an integer")
        if not isinstance(persona["pain_points"], list) or not all(isinstance(p, str) for p in persona["pain_points"]):
            raise ValueError("pain_points must be a list of strings")
        if not isinstance(persona["motivations"], list) or not all(isinstance(p, str) for p in persona["motivations"]):
            raise ValueError("motivations must be a list of strings")
        if not isinstance(persona["scenarios"], list) or not all(isinstance(s, str) for s in persona["scenarios"]):
            raise ValueError("scenarios must be a list of strings")
        if not isinstance(persona["demographics"], dict) or not all(k in persona["demographics"] for k in ["location", "family_status", "workplace"]):
            raise ValueError("demographics must be a dict with location, family_status, workplace")
        # Normalize adoption_stage
        adoption_stages = ["Innovator", "Early Adopter", "Early Majority", "Late Majority", "Laggard"]
        adoption_stage = persona["adoption_stage"].title()
        if adoption_stage not in adoption_stages:
            raise ValueError(f"Invalid adoption stage: {persona['adoption_stage']}")
        persona["adoption_stage"] = adoption_stage
        # Fix gender inconsistencies
        family_status = persona["demographics"]["family_status"].lower()
        if "mother" in family_status and persona["gender"].lower() != "female":
            persona["gender"] = "Female"
        elif "father" in family_status and persona["gender"].lower() != "male":
            persona["gender"] = "Male"
        # Clean family_status
        persona["demographics"]["family_status"] = re.sub(r'^(she|he|they)\s+', '', family_status, flags=re.IGNORECASE).strip().title()
        logger.info(f"Validated persona: {persona['name']} ({persona['adoption_stage']})")
        return persona
    except json.JSONDecodeError as e:
        logger.error(f"JSON parsing error: {str(e)}")
        return None
    except Exception as e:
        logger.error(f"Invalid persona JSON: {str(e)}")
        return None

def setup_agent(vectorstore):
    retriever_tool = create_retriever_tool(
        vectorstore.as_retriever(search_kwargs={"k": 5}),
        "course_materials_retriever",
        "Searches course materials on optimization, simulation, decision tech, project management, and lean startup."
    )
    web_tool = TavilySearch(max_results=3)
    tools = [retriever_tool, web_tool, scenario_simulator]
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
    prompt = ChatPromptTemplate.from_messages([
        ("system", """
        You are Dr. Sue Cohen, a passionate professor of Commercializing New Technologies, unless a persona is active, in which case you respond fully in-character as the persona. Answer in a witty, professional, detailed style (150-250 words for course content to ensure depth with examples and in-depth analysis), using main points from diverse course materials without citing file names in the main response. Avoid over-relying on any single lecture (e.g., Lecture 5). Use Tavily/web for recent updates in separate paragraphs with up to 5 URLs only if relevant and available; if not, omit the web section. Don’t mix course and web content. Run scenario sims for adoption/scalability with clear params. If a persona is active, respond only as that persona, staying true to their profile. Redirect off-topic queries, prompt next steps with a sharp question.
        Active Persona (if any): {active_persona}
        """),
        ("human", "{input}"),
        MessagesPlaceholder("agent_scratchpad")
    ])
    agent = create_tool_calling_agent(llm, tools, prompt)
    executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    return executor

def chat(message, history):
    global active_persona, conversation_count
    logger.info(f"Received message: {message}")
    
    # Parse Persona Mode commands
    if message.lower().startswith("persona mode:"):
        command = message[len("persona mode:"):].strip().lower()
        logger.info(f"Processing persona command: {command}")
        
        if command.startswith("generate persona:"):
            try:
                user_input = command[len("generate persona:"):].strip()
                logger.info(f"Processing generate persona with input: {user_input}")
                personas = persona_generator(user_input)
                logger.info(f"Generated personas: {[p['id'] + ': ' + p['name'] for p in personas]}")
                return json.dumps(personas, indent=2) + "\n\nCopy the JSON above and save to a .json file for reuse. Use 'Persona Mode: Activate Persona: [JSON]' to interact with a persona."
            except Exception as e:
                logger.error(f"Persona generation failed in chat: {e}")
                return f"Error: {e}. Use format: 'Persona Mode: Generate Persona: [description]'"
        
        elif command.startswith("activate persona:"):
            try:
                persona_json = command[len("activate persona:"):].strip()
                persona = validate_persona_json(persona_json)
                if persona:
                    active_persona = persona
                    conversation_count = 0
                    # Set pronouns based on gender
                    pronouns = {
                        "Male": {"subject": "he", "possessive": "his"},
                        "Female": {"subject": "she", "possessive": "her"},
                        "Non-binary": {"subject": "they", "possessive": "their"}
                    }.get(persona["gender"].title(), {"subject": "they", "possessive": "their"})
                    # Clean fields for summary
                    pain_point_1 = re.sub(r'^(she|he|they)\s+', '', persona['pain_points'][0], flags=re.IGNORECASE).strip().lower()
                    pain_point_2 = re.sub(r'^(she|he|they)\s+', '', persona['pain_points'][1], flags=re.IGNORECASE).strip().lower() if len(persona['pain_points']) > 1 else "workplace challenges"
                    motivation_1 = re.sub(r'^(she|he|they)\s+', '', persona['motivations'][0], flags=re.IGNORECASE).strip().lower()
                    motivation_2 = re.sub(r'^(she|he|they)\s+', '', persona['motivations'][1], flags=re.IGNORECASE).strip().lower() if len(persona['motivations']) > 1 else "professional growth"
                    scenario_1 = re.sub(r'^(she|he|they)\s+', '', persona['scenarios'][0], flags=re.IGNORECASE).strip()
                    scenario_2 = re.sub(r'^(she|he|they)\s+', '', persona['scenarios'][1], flags=re.IGNORECASE).strip() if len(persona['scenarios']) > 1 else f"Evaluates tech to support {motivation_1}."
                    # Generate summary in Dr. Sue's voice
                    summary = f"""
{persona['name']} is a quintessential {persona['adoption_stage'].lower()} in the wild world of healthcare commercialization, tackling innovation at {persona['demographics']['workplace'].lower()}. {pronouns['subject'].title()}'s passion for {motivation_1} drives {pronouns['subject']} to size up healthcare technologies, but {pain_point_1} keeps {pronouns['subject']} on edge. {scenario_1} Meanwhile, {pronouns['subject']} grapples with {pain_point_2}, a stark reminder of the stakes in {persona['demographics']['location'].lower()}. Yet, {pronouns['subject']}'s commitment to {motivation_2} fuels {pronouns['possessive']} resolve, juggling {persona['demographics']['family_status'].lower()} alongside {pronouns['possessive']} career. In our course's spirit, {pronouns['subject']}'s a case study in navigating adoption barriers—think regulatory hoops or staff pushback—demanding a pitch sharper than a scalpel. To win {pronouns['subject']} over, you'll need rock-solid data and a plan to scale without derailing {pronouns['possessive']} hectic life. What's your strategy to crack {pronouns['possessive']} skepticism?
                    """
                    logger.info(f"Activated persona: {persona['name']} ({persona['adoption_stage']})")
                    return f"Activated persona: {persona['name']} ({persona['adoption_stage']}). Ask up to {MAX_QUESTIONS} questions.\n\n{summary.strip()}"
                return f"Error: Invalid persona JSON. Ensure all required fields are included and JSON is valid. Example: {{'id': 'uuid', 'name': 'Name', 'age': 35, 'gender': 'Male', 'occupation': 'Doctor', 'adoption_stage': 'Innovator', 'pain_points': ['issue1', 'issue2'], 'motivations': ['goal1', 'goal2'], 'demographics': {{'location': 'City', 'family_status': 'Status', 'workplace': 'Place'}}, 'scenarios': ['Story1', 'Story2']}}"
            except Exception as e:
                logger.error(f"Persona activation failed: {e}")
                return f"Error: {e}. Use format: 'Persona Mode: Activate Persona: [JSON]'"
        
        elif command == "end conversation":
            if active_persona:
                name = active_persona['name']
                active_persona = None
                conversation_count = 0
                logger.info(f"Ended conversation with {name}")
                return f"Ended conversation with {name}. Now in Dr. Sue Cohen's TA mode. What's your next question?"
            logger.warning("No active persona conversation to end")
            return "No active persona conversation to end."
        
        else:
            logger.error(f"Invalid persona command: {command}")
            return f"Invalid Persona Mode command: {command}. Use 'Generate Persona: [description]', 'Activate Persona: [JSON]', or 'End Conversation'."

    # Parse scenario simulator query
    if "initial_users" in message.lower():
        params = {}
        try:
            for pair in message.split(","):
                key, value = pair.split(":")
                key = key.strip()
                value = float(value.strip().replace("%", "")) / 100 if "%" in value else int(value.strip())
                params[key] = value
            logger.info(f"Running scenario simulator with params: {params}")
            return scenario_simulator(params)
        except Exception as e:
            logger.error(f"Failed to parse scenario params: {e}")
            return f"Error parsing params: {e}. Format: initial_users:1000, growth_rate:5%, time_months:12, uncertainty:10%"

    # Handle persona interaction
    if active_persona and conversation_count < MAX_QUESTIONS:
        try:
            llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
            messages = [
                ("system", f"""
                You are {active_persona['name']}, a {active_persona['age']}-year-old {active_persona['gender']} {active_persona['occupation']}.
                Adoption stage: {active_persona['adoption_stage']}.
                Personality traits: {active_persona['adoption_stage'].lower()} (trailblazers for Innovators, visionaries for Early Adopters, pragmatic for Early Majority, skeptical for Late Majority, change-resistant for Laggards).
                Pain points: {', '.join(active_persona['pain_points'])}.
                Motivations: {', '.join(active_persona['motivations'])}.
                Demographics: {json.dumps(active_persona['demographics'])}.
                You are being interviewed about a healthcare technology (e.g., AI patient monitoring tool) as a potential customer or stakeholder. Respond in character, using first-person voice, in 3-5 detailed sentences. Draw on your pain points (e.g., regulatory barriers, staff resistance from course materials), motivations (e.g., evidence-based validation for scaling), demographics (e.g., work-life balance in Pittsburgh), and scenarios (e.g., overcoming adoption hurdles like IT integration). Use conversational, realistic language with specific examples from healthcare commercialization (e.g., pilot studies, board approvals), avoiding medical jargon unless relevant to your role. Do NOT ask a question back to the interviewer; provide a complete, standalone response to simulate a realistic client interview for student training.
                Question: {message}
                """),
                ("human", message)
            ]
            response = llm.invoke(messages)
            conversation_count += 1
            remaining = MAX_QUESTIONS - conversation_count
            logger.info(f"Persona response generated for {active_persona['name']}, {remaining} questions left")
            return f"{response.content}\n\n({remaining} questions left. Use 'Persona Mode: End Conversation' to stop.)"
        except Exception as e:
            logger.error(f"Persona interaction failed: {e}")
            return f"Error: {e}. Try again or end conversation."

    # Handle regular RAG queries
    retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
    docs = retriever.invoke(message)
    context = "\n\n".join([doc.page_content for doc in docs])
    titles = [os.path.basename(doc.metadata.get('source', 'Unknown')) for doc in docs]
    logger.info(f"Retrieved {len(docs)} documents for query: {message}")
    
    try:
        web_results = TavilySearch(max_results=5).invoke(message)
        if isinstance(web_results, str):
            web_context = f"Web: {web_results[:150]}..."
        else:
            web_context = "\n".join([f"Web: {res.get('title', 'Untitled')} - {res.get('content', '')[:150]}... (URL: {res.get('url', 'Unknown')})" for res in web_results])
        logger.info(f"Retrieved {len(web_results) if isinstance(web_results, list) else 1} web results from Tavily")
    except Exception as e:
        web_context = ""
        logger.error(f"Tavily search failed: {e}")
    
    input_with_context = f"{message}\n\nCourse Materials:\n{context}\n\nWeb Updates:\n{web_context}\n\nPDF References:\n{', '.join(titles)}"
    logger.info(f"Processing RAG query: {message}")
    return agent.invoke({"input": input_with_context, "active_persona": active_persona or "None"})["output"]

if not os.path.exists("faiss_index"):
    logger.warning("FAISS index not found, rebuilding...")
    vectorstore = build_vectorstore()
else:
    logger.info("Loading prebuilt FAISS index...")
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    vectorstore = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)

agent = setup_agent(vectorstore)

demo = gr.ChatInterface(
    fn=chat,
    chatbot=gr.Chatbot(type='messages'),
    title="COMMERCIALIZING NEW TECHNOLOGIES RAG Bot",
    description="Query course materials, run scenario simulations, or use Persona Mode: 'Generate Persona: [description]', 'Activate Persona: [JSON]', 'End Conversation'"
)

if __name__ == "__main__":
    logger.info("Starting Gradio app...")
    try:
        
        logger.info(f"Using port: {7860}")
        demo.launch(server_name="0.0.0.0", server_port=7860)
    except Exception as e:
        logger.error(f"Failed to start Gradio app: {e}")
        raise