import asyncio
import os
import sys

# Add current directory to path and import prompts
current_dir = os.path.dirname(__file__)
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

import prompts
from IPython.display import display, Markdown
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.models.lite_llm import LiteLlm
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.genai import types
from google.adk.tools.langchain_tool import LangchainTool
from langchain_tavily import TavilySearch

load_dotenv()

APP_NAME = "startup_validator"
USER_ID = "arindam_1729"
SESSION_ID = "startup_validation_session"


async def run_validation(idea: str, nebius_api_key: str, tavily_api_key: str):
    if not nebius_api_key:
        raise ValueError("Missing Nebius API key")
    if not tavily_api_key:
        raise ValueError("Missing Tavily API key")

    # Build model and tools using provided API keys
    nebius_llm = LiteLlm(
        model="nebius/Qwen/Qwen3-235B-A22B-Instruct-2507", api_key=nebius_api_key
    )

    tavily_tool_instance = TavilySearch(
        max_results=3,
        search_depth="basic",
        include_answer=True,
        include_raw_content=False,
        include_images=False,
        tavily_api_key=tavily_api_key,
    )
    tavily_search = LangchainTool(tool=tavily_tool_instance)

    # Define agents
    idea_clarifier_agent = LlmAgent(
        name="IdeaClarifierAgent",
        model=nebius_llm,
        instruction=prompts.IDEA_PROMPT,
        description="Helps clarify and refine the startup idea.",
        output_key="clarified_idea",
    )

    market_research_agent = LlmAgent(
        name="MarketResearchAgent",
        model=nebius_llm,
        instruction=prompts.MARKET_RESEARCH_PROMPT,
        description="Conducts market research for the startup idea.",
        tools=[tavily_search],
        output_key="market_research",
    )

    competitor_analysis_agent = LlmAgent(
        name="CompetitorAnalysisAgent",
        model=nebius_llm,
        instruction=prompts.COMPETITOR_ANALYSIS_PROMPT,
        description="Conducts competitor analysis for the startup idea.",
        tools=[tavily_search],
        output_key="competitor_analysis",
    )

    report_agent = LlmAgent(
        name="ReportAgent",
        model=nebius_llm,
        instruction=prompts.REPORT_PROMPT,
        description="Generates a report based on the analysis findings.",
        output_key="validation_report",
    )

    startup_validation_agent = SequentialAgent(
        name="StartupValidationAgent",
        sub_agents=[
            idea_clarifier_agent,
            market_research_agent,
            competitor_analysis_agent,
            report_agent,
        ],
        description="Validates startup ideas through a structured analysis process.",
    )

    initial_state = {"idea": idea}
    session_service = InMemorySessionService()
    await session_service.create_session(
        app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID, state=initial_state
    )

    runner = Runner(
        agent=startup_validation_agent,
        app_name=APP_NAME,
        session_service=session_service,
    )

    content = types.Content(role="user", parts=[types.Part(text=idea)])
    
    # Handle the runner.run() method properly
    run_result = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content)
    
    # Check if it's a generator and consume it
    if hasattr(run_result, '__aiter__'):
        # It's an async generator, consume it
        async for _ in run_result:
            pass  # Just consume the generator
    elif hasattr(run_result, '__await__'):
        # It's a coroutine, await it normally
        await run_result
    else:
        # It's a regular generator, consume it
        try:
            while True:
                next(run_result)
        except StopIteration:
            pass

    session = await session_service.get_session(
        app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
    )
    
    # Debug: Print session state for troubleshooting
    print(f"Session state keys: {list(session.state.keys())}")
    print(f"Session state: {session.state}")

    import ast

    def safe_parse(val):
        if isinstance(val, dict):
            return val
        try:
            return ast.literal_eval(val)
        except Exception:
            return {}

    market_research = safe_parse(session.state.get("market_research", {}))
    competitor_analysis = safe_parse(session.state.get("competitor_analysis", {}))
    validation_report = safe_parse(session.state.get("validation_report", {}))

    summary = f"""
🎉 **STARTUP IDEA VALIDATION COMPLETED!**

## 📊 Validation Summary

- **Startup Idea:** {idea}
- **Idea Clarification:** ✅ Completed
- **Market Research:** ✅ Completed
- **Competitor Analysis:** ✅ Completed
- **Final Report:** ✅ Generated

## 📈 Key Market Insights

### **TAM:** 
{market_research.get("total_addressable_market", "")}
### **Target Segments:** 
{market_research.get("target_customer_segments", "")}
### **SOM:** 
{market_research.get("serviceable_obtainable_market", "")}
### **SAM:** 
{market_research.get("serviceable_available_market", "")}

## 🏆 Competitive Positioning

{competitor_analysis.get("competitors", "")}

{competitor_analysis.get("swot_analysis", "")}

{competitor_analysis.get("positioning", "")}

---

## 📋 Comprehensive Validation Report

{validation_report.get("executive_summary", "")}

{validation_report.get("idea_assessment", "")}

{validation_report.get("market_opportunity", "")}

{validation_report.get("competitive_landscape", "")}

{validation_report.get("recommendations", "")}

{validation_report.get("next_steps", "")}

---

> ⚠️ *Disclaimer: This validation is for informational purposes only. Conduct additional due diligence before making investment decisions.*

"""

    return summary


if __name__ == "__main__":
    nebius_key = os.getenv("NEBIUS_API_KEY", "")
    tavily_key = os.getenv("TAVILY_API_KEY", "")
    asyncio.run(
        run_validation(
            "A CodeReview Agent that reviews your code in each PR",
            nebius_key,
            tavily_key,
        )
    )