import os
import requests
from utils.prompts import LPA_STATEMENT_PROMPT, SIDE_LETTER_PROMPT
from backendAPIs import load_partner_data
from llm_bedrock import model
import re
import json
from datetime import datetime, timedelta
from typing import List, Dict
import json


def generate_quarterly_management_fees(
        fund_start_date: str,
        fund_end_date: str,
        fee_schedule: List[Dict[str, float]]
) -> str:
    """
    Generates a JSON string of quarterly management fees based on the provided schedule.

    :param fund_start_date: Fund start date in 'YYYY-MM-DD' format.
    :param fund_end_date: Fund end date in 'YYYY-MM-DD' format.
    :param fee_schedule: List of fee changes with 'start_date' and 'annual_fee_percent'.
    :return: JSON string representing quarterly management fees.
    """
    # Convert string dates to datetime objects
    start_date = datetime.strptime(fund_start_date, "%Y-%m-%d")
    end_date = datetime.strptime(fund_end_date, "%Y-%m-%d")

    # Sort fee_schedule by start_date
    fee_schedule_sorted = sorted(fee_schedule, key=lambda x: x['start_date'])

    # Create list of fee periods with start and end dates
    fee_periods = []
    for i, fee in enumerate(fee_schedule_sorted):
        period_start = datetime.strptime(fee['start_date'], "%Y-%m-%d")
        if i + 1 < len(fee_schedule_sorted):
            period_end = datetime.strptime(fee_schedule_sorted[i + 1]['start_date'], "%Y-%m-%d") - timedelta(days=1)
        else:
            period_end = end_date
        # Ensure period doesn't start before fund_start_date
        if period_start < start_date:
            period_start = start_date
        # Ensure period_end doesn't exceed fund_end_date
        if period_end > end_date:
            period_end = end_date
        fee_periods.append({
            'start_date': period_start,
            'end_date': period_end,
            'annual_management_fee_percent': fee['annual_management_fee_percent']
        })

    # Function to get quarter for a given date
    def get_quarter(date):
        return (date.month - 1) // 3 + 1

    # Initialize a dictionary to hold yearly fees
    yearly_fees = {}

    current_date = start_date
    while current_date <= end_date:
        year = current_date.year
        if year not in yearly_fees:
            yearly_fees[year] = [0.0, 0.0, 0.0, 0.0]

        quarter = get_quarter(current_date)

        # Find the applicable fee for the current_date
        applicable_fee = 0.0
        for period in fee_periods:
            if period['start_date'] <= current_date <= period['end_date']:
                applicable_fee = round(period['annual_management_fee_percent'] / 4, 6)  # Convert to decimal per quarter
                break

        yearly_fees[year][quarter - 1] = applicable_fee

        # Move to the next quarter
        if quarter == 4:
            # Next quarter is Q1 of next year
            current_date = datetime(year + 1, 1, 1)
        else:
            # Next quarter: first month of next quarter
            next_quarter_month = 3 * quarter + 1
            current_date = datetime(year, next_quarter_month, 1)

    # Fill in zero fees before start_date and after end_date
    formatted_yearly_fees = {}
    for year in range(start_date.year, end_date.year + 1):
        if year in yearly_fees:
            formatted_yearly_fees[str(year)] = yearly_fees[year]
        else:
            formatted_yearly_fees[str(year)] = [0.0, 0.0, 0.0, 0.0]

    # Convert to JSON string with keys as strings
    json_output = json.dumps(formatted_yearly_fees, indent=4)
    return json_output



def lpa_amendment(file_data):


    FINAL_LPA_PROMPT = LPA_STATEMENT_PROMPT.format(document=file_data)

    json_string = model.invoke(FINAL_LPA_PROMPT).content

    pattern = r'<output>(.*?)</output>'
    match = re.search(pattern, json_string, re.DOTALL)

    try:
        if match:
          final_json_data = match.group(1).strip()

          lpa_output = json.loads(final_json_data)

          if 'managementFee' in lpa_output and all(key in lpa_output['managementFee'] for key in
                                             ['fund_start_date', 'fund_end_date',
                                              'management_fee_schedule']):
              formatted_management_fee= generate_quarterly_management_fees( lpa_output['managementFee']['fund_start_date'],lpa_output['managementFee']['fund_end_date'],lpa_output['managementFee']['management_fee_schedule'])

              lpa_output['managementFee'] = formatted_management_fee

              final_json_data = json.dumps(lpa_output)

              print("Final json")

              print(final_json_data)

        return  final_json_data
    except Exception as e:
        print("")
        return None


def create_xml_string(data):
    xml_string = '<partners>\n'
    for item in data:
        xml_string += f'  <partner>\n'
        xml_string += f'    <id>{item.get("_id", "")}</id>\n'
        xml_string += f'    <partnername>{item.get("partnerExactName", "Not specified")}</partnername>\n'
        xml_string += f'  </partner>\n'
    xml_string += '</partners>'
    return xml_string


def side_letter_doc(type,entity_id,file_data,role_data):

    output = model.invoke(SIDE_LETTER_PROMPT.format(document=file_data,role_doc=create_xml_string(role_data))).content

    pattern = r'<output>(.*?)</output>'
    match = re.search(pattern, output, re.DOTALL)

    json_data = None
    role_id = None

    if match:
        json_data = match.group(1).strip()
        ## get id
        pattern_id =  r'<id>(.*?)</id>'
        match = re.search(pattern_id, output, re.DOTALL)
        if match:
            role_id = match.group(1).strip()

    try:

        lpa_output = json.loads(json_data)

        if 'managementFee' in lpa_output and all(key in lpa_output['managementFee'] for key in
                                                     ['fund_start_date', 'fund_end_date',
                                                      'management_fee_schedule']):
            formatted_management_fee = generate_quarterly_management_fees(
                    lpa_output['managementFee']['fund_start_date'], lpa_output['managementFee']['fund_end_date'],
                    lpa_output['managementFee']['management_fee_schedule'])

            lpa_output['managementFee'] = formatted_management_fee

        final_json_data = json.dumps(lpa_output)

        print("Final json")

        print(final_json_data)

    except Exception as e:
        print("")

    if role_id is None:
        return "Role is Not fund"

    response = load_partner_data(type,entity_id,role_id,final_json_data)


    return response , role_id