import json
import ast
import os
import datetime
from openai import AsyncOpenAI
from edgar import set_identity, get_filings, get_by_accession_number, Company
from chainlit.playground.providers.openai import stringify_function_call
import chainlit as cl
from dotenv import load_dotenv

load_dotenv()

openai_api_key = os.getenv("OPENAI_API_KEY")
edgar_profile = os.getenv("SEC_EDGAR_PROFILE")

print(openai_api_key)
print(edgar_profile)

set_identity(edgar_profile)

client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))

MAX_ITER = 10

cl.instrument_openai()


def get_latest_filings(year: int = 2024, quarter: int = 1, form: str = "NPORT-P"):
    """
    Get the latest filings
    """

    filings = get_filings(year, quarter, form=form)
    return str(filings)


def get_filing_by_accession_number(accession_number: str):
    """
    Get the filing by accession number
    """

    filing = get_by_accession_number(accession_number)

    return str(filing)


def get_financials_facts_for_company(ticker_or_CIK: str):
    """
    Get the financials for a company
    """

    company = Company(ticker_or_CIK)
    return str(company.get_facts())


def get_financials_from_filing(accession_number: str):
    """
    Get the financials from a filing that is taken from a TenK or TenQ
    """

    filing = get_by_accession_number(accession_number).obj()
    financials = filing.financials
    financials.balance_sheet
    return str(financials.balance_sheet)


def get_press_releases(ticker_or_CIK: str, num_filings: int = 5):
    """
    Get the press releases from latest 8Ks
    """

    company = Company(ticker_or_CIK)
    latest = company.get_filings(form="8-K")

    eightk = [x.obj() for x in latest.latest(num_filings)]

    with_press_releases = [x for x in eightk if x.has_press_release]

    press_releases = [x.press_releases for x in with_press_releases]

    urls = [x.attachments.get(0).url for x in press_releases]

    return str(urls)


tools = [
    {
        "type": "function",
        "function": {
            "name": "get_latest_filings",
            "description": "Returns the latest filings from SEC. All parameters are optionals, default are year 2024, quarter 1 and form NPORT-P.",
            "parameters": {
                "type": "object",
                "properties": {
                    "year": {
                        "type": "number",
                        "description": "The ",
                    },
                    "form": {
                        "type": "string",
                        "enum": ["NPORT-P", "10-Q", "10-K", "8-K"],
                    },
                    "quarter": {"type": "number", "enum": [1, 2, 3, 4]},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_filing_by_accession_number",
            "description": "Returns the content of a SEC filing in markdown format by the given accession number.",
            "parameters": {
                "type": "object",
                "properties": {
                    "accession_number": {
                        "type": "string",
                        "description": "The accession number of the SEC filing to retrieve.",
                    }
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_financials_facts_for_company",
            "description": "Returns financial facts for a company identified by its ticker symbol or CIK (Central Index Key).",
            "parameters": {
                "type": "object",
                "properties": {
                    "ticker_or_CIK": {
                        "type": "string",
                        "description": "The ticker symbol or CIK of the company for which to retrieve financial facts.",
                    }
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_financials_from_filing",
            "description": "Returns financial from a filing that is taken from a TenK or TenQ. Stockholders' Equity, Liabilities and Assets can be taken using this function.",
            "parameters": {
                "type": "object",
                "properties": {
                    "accession_number": {
                        "type": "string",
                        "description": "The accession number of the SEC filing to retrieve financial status.",
                    }
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_press_releases",
            "description": "Returns the press releases from the latest 8-K filings of a company identified by its ticker symbol or CIK.",
            "parameters": {
                "type": "object",
                "properties": {
                    "ticker_or_CIK": {
                        "type": "string",
                        "description": "The ticker symbol or CIK of the company for which to retrieve press releases.",
                    },
                    "num_filings": {
                        "type": "number",
                        "description": "The number of 8-K filings to consider for press releases. Defaults to 5 which is enough.",
                    },
                },
            },
        },
    },
]

prompt = f"""
You are a helpful SEC/Edgar financial assistant. Your job is to use the tools you have available to help
the user with their financial queries. You have access to the following tools:

- get_latest_filings
- get_filing_by_accession_number
- get_financials_facts_for_company
- get_financials_from_filing
- get_press_releases

You can combine them and plan your responses accordingly.

For your information, the current date is:
{datetime.datetime.now().strftime("%m-%d-%Y")}
"""

print(prompt)


@cl.on_chat_start
def start_chat():
    cl.user_session.set(
        "message_history",
        [
            {
                "role": "system",
                "content": prompt,
            }
        ],
    )


@cl.step(type="tool")
async def call_tool(tool_call_id, name, arguments, message_history):
    arguments = ast.literal_eval(arguments)

    current_step = cl.context.current_step
    current_step.name = name
    current_step.input = arguments

    current_step.input = arguments
    if name == "get_latest_filings":
        function_response = get_latest_filings(**arguments)
    elif name == "get_filing_by_accession_number":
        function_response = get_filing_by_accession_number(**arguments)
    elif name == "get_financials_facts_for_company":
        function_response = get_financials_facts_for_company(**arguments)
    elif name == "get_press_releases":
        function_response = get_press_releases(**arguments)
    elif name == "get_financials_from_filing":
        function_response = get_financials_from_filing(**arguments)

    current_step.output = function_response
    current_step.language = "json"

    message_history.append(
        {
            "role": "function",
            "name": name,
            "content": function_response,
            "tool_call_id": tool_call_id,
        }
    )


@cl.step(type="Financial AI")
async def call_gpt4(message_history):
    settings = {
        "model": "gpt-4o",
        "tools": tools,
        "tool_choice": "auto",
        "temperature": 0,
    }

    cl.context.current_step.generation = cl.ChatGeneration(
        provider="openai-chat",
        messages=[
            cl.GenerationMessage(
                formatted=m["content"], name=m.get("name"), role=m["role"]
            )
            for m in message_history
        ],
        settings=settings,
    )

    stream = await client.chat.completions.create(
        messages=message_history, stream=True, **settings
    )

    tool_call_id = None
    function_output = {"name": "", "arguments": ""}

    final_answer = cl.Message(content="", author="Answer")

    async for part in stream:
        new_delta = part.choices[0].delta
        tool_call = new_delta.tool_calls and new_delta.tool_calls[0]
        function = tool_call and tool_call.function
        if tool_call and tool_call.id:
            tool_call_id = tool_call.id

        if function:
            if function.name:
                function_output["name"] = function.name
            else:
                function_output["arguments"] += function.arguments
            await cl.context.current_step.stream_token(
                json.dumps(function_output), is_sequence=True
            )
        if new_delta.content:
            await cl.context.current_step.stream_token(new_delta.content)
            if not final_answer.content:
                await final_answer.send()
            await final_answer.stream_token(new_delta.content)

    cl.context.current_step.generation.completion = cl.context.current_step.output

    if tool_call_id:
        cl.context.current_step.output = stringify_function_call(function_output)
        cl.context.current_step.language = "json"
        await call_tool(
            tool_call_id,
            function_output["name"],
            function_output["arguments"],
            message_history,
        )

    if final_answer.content:
        await final_answer.update()

    return tool_call_id


@cl.on_message
async def on_message(message: cl.Message):
    message_history = cl.user_session.get("message_history")
    message_history.append({"role": "user", "content": message.content})

    cur_iter = 0

    while cur_iter < MAX_ITER:
        tool_call_id = await call_gpt4(message_history)
        if not tool_call_id:
            break

        cur_iter += 1