# Import required modules
import requests
from bs4 import BeautifulSoup
import nest_asyncio
import asyncio
import json
import re
from crawl4ai import *
import os
from dotenv import load_dotenv
import google.generativeai as genai

# Load environment variables from a .env file
load_dotenv()  # Make sure a .env file exists with GOOGLE_API_KEY=<your_api_key>
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")  # Fetch the API key

# Apply nest_asyncio to enable asynchronous tasks in Jupyter/interactive environments
nest_asyncio.apply()

# Asynchronous function to extract text from a website
async def extract_text_from_website(url):
    async with AsyncWebCrawler() as crawler:
        result = await crawler.arun(url=url)
        return result.markdown


# Define market sectors

# Define the prompt for generating market scenarios
# Configure the generative AI model
genai.configure(api_key=GOOGLE_API_KEY)  # Replace with your API key

generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-2.0-flash-exp",
    generation_config=generation_config,
)

chat_session = model.start_chat()

# Function to get a response from the generative AI model
def get_response(llm, prompt):
    response = llm.send_message(prompt)
    return response



# Function to extract JSON content from the response
def extract_json_content(text):
    match = re.search(r"```json\n(.*?)```", text, re.DOTALL)
    if match:
        return match.group(1).strip()
    else:
        return None


if __name__ == "__main__":
    # Extract market data from the given URL
    url = "https://www.livemint.com/market/stock-market-news/page-7"
    context_data = asyncio.run(extract_text_from_website(url))


    sectors = [
        "Communication Services",
        "Consumer Discretionary",
        "Consumer Staples",
        "Energy",
        "Financials",
        "Health Care",
        "Industrials",
        "Information Technology",
        "Materials",
        "Real Estate",
        "Utilities",
    ]

    prompt = f"""
    # TASK: Analyze market context and identify potential market scenarios.

    # CONTEXT:
    {context_data}
    # END CONTEXT

    # INSTRUCTION: Based on the provided market context, analyze and identify up to three plausible market scenarios.
    # For each scenario, determine its name (e.g., "Moderate Downturn"), the general market direction ("up" or "down"), a major trigger point that could cause the scenario to unfold, and a list of sectors that would be significantly impacted. Each 'sector_impact' list should have less than or equal to 4 sectors.

    # OUTPUT FORMAT: Provide the analysis in JSON format with the following structure.
    # Use the sector names provided:
    {sectors}

    # EXAMPLE:
    ```json
    {{
        "market_scenarios": {{
            "scenario1": {{
                "name": "Moderate Downturn",
                "direction": "down",
                "trigger": "Interest rate hike",
                "sector_impact": [
                    "Financials",
                    "Energy"
                ]
            }},
            "scenario2": {{
                "name": "Bullish Growth",
                "direction": "up",
                "trigger": "Successful vaccine rollout",
                "sector_impact": [
                    "Health Care",
                    "Information Technology"
                ]
            }}
        }}
    }}
    """



    # Generate the response
    answer = get_response(chat_session, prompt)

    # Extract the JSON output from the response
    json_output = extract_json_content(answer.text)

    # Define output file path
    output_file = "output_files/scenario.json"

    # Parse the output into a JSON object and save it to a file
    try:
        analysis_json = json.loads(json_output)
        os.makedirs(os.path.dirname(output_file), exist_ok=True)  # Ensure the output directory exists
        with open(output_file, "w") as f:
            json.dump(analysis_json, f, indent=4)  # Save JSON to a file with indentation
        print(f"Analysis saved to '{output_file}'")
    except json.JSONDecodeError:
        print("Error: Could not decode the output from the model into JSON format.")
    except Exception as e:
        print(f"Error: {e}")