from datetime import datetime
import re
import os
import json

import click
import yaml

from jinja2 import Template
from jinja2 import Environment, BaseLoader
from openai import OpenAI
import logging
from maiden import __version__

logging.basicConfig(level=logging.INFO)

def load_data(file_path):
    """Load data from a YAML or JSON file"""
    if file_path.endswith('.yaml') or file_path.endswith('.yml'):
        with open(file_path, 'r') as f:
            return yaml.safe_load(f)
    elif file_path.endswith('.json'):
        with open(file_path, 'r') as f:
            return json.load(f)
    else:
        with open(file_path, 'r') as f:
            content = f.read()
            if "END %>" in content:
                content = content.replace("END %>", "END \%>")
            return { "content" : content }

def openai_prompt(prompt):
    """Generate text from OpenAI based on the prompt"""
    api_key = os.environ["OPENAI_API_KEY"]
    client = OpenAI(api_key=api_key)

    messages = [
        {"role": "system", "content": "You are a helpful assistant for writing techinal documents"},
        {"role": "user", "content": f"Input: {prompt}"},
        {"role": "user", "content": "Please provide the requested information below:"}
    ]

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        # max_tokens=2048,
        stop=None,
        temperature=0.1,
        frequency_penalty=0.5,
        presence_penalty=0.5,
    )

    corrected_text = response.choices[0].message.content.strip()
    return corrected_text.strip()

def postprocess_template(template_str):
    """Preprocess the template to replace custom syntax with actual data"""
    pattern = r"<% FROM (.+?) GET '(.+?)' %>"
    matches = re.findall(pattern, template_str)

    for match in matches:
        source_file, key = match
        data = load_data(source_file)
        
        # Split the key by dots to traverse nested dictionaries
        keys = key.split('.')
        value = data
        for key_ in keys:
            if isinstance(value, list):
                filtered_list = [item for item in value if isinstance(item, dict) and key_ in item]
                logging.info(f"List value = {filtered_list}")
                if filtered_list:
                    value = filtered_list[0]
                else:
                    value = '[Missing Data]'
                    break
            elif isinstance(value, dict):
                value = value.get(key_, '[Missing Data]')
                if value == '[Missing Data]':
                    break
        
        # Handle cases where the key path is invalid
        if not value:
            value = "[Missing Data]"

        # Replace the custom directive with the actual value
        template_str = template_str.replace(f"<% FROM {source_file} GET '{key}' %>", f"{value}")
    
    date_pattern = r"<% DATE (.+?) %>"
    matches = re.findall(date_pattern, template_str)

    for match in matches:
        date_format = match.strip()
        value = datetime.now().strftime(date_format)
        template_str = template_str.replace(f"<% DATE {date_format} %>", f"{value}")

    # Handle OpenAI prompts
    openai_pattern = r"<% OPENAI(.+?)END %>"
    matches = re.findall(openai_pattern, template_str, re.DOTALL)

    for match in matches:
        prompt_template = match.strip()
        logging.info(f"OPENAI PART {prompt_template}")
        value = openai_prompt(prompt_template)
        logging.info(f"VALUE = {value}")
        template_str = template_str.replace(f"<% OPENAI{match}END %>", value)

    return template_str

def process_template(data, processed_template_str):
    def file_exists(path):
        return os.path.isfile(path)
    
    # Create an environment
    env = Environment(loader=BaseLoader())
    
    # Register the custom filter
    env.filters['file_exists'] = file_exists
    
    # Create a template from the string
    template = env.from_string(processed_template_str)
    
    # Render the template with the given data
    return template.render(data)

def process_data_template_file(data_file):
    logging.info("Processing Data Template File")
    # Load the template from the file
    with open(data_file) as f:
        template_str = f.read()
    
    processed_template_str = postprocess_template(template_str)
    return yaml.safe_load(processed_template_str)

@click.command()
@click.argument('template_file')
@click.argument('output_file')
@click.option('--data-file', '-d', help='Data files in YAML format')
@click.option('--process-data-template', is_flag=True, default=True, help="Process data file as a template")
@click.version_option(__version__)
def generate_md(template_file, output_file, data_file, process_data_template):
    """CLI command to generate markdown from template"""
    if data_file:
        if process_data_template:
            data = process_data_template_file(data_file)
        else:
            data = load_data(data_file)
    
        logging.info(f"Data = {data}")

    with open(template_file, 'r') as f:
        template_str = f.read()

    # Preprocess template to replace custom syntax
    if data_file:
        template_str = process_template(data, template_str)
    template_str = postprocess_template(template_str)
    
    # Write output to file
    with open(output_file, 'w') as f:
        f.write(template_str)

if __name__ == '__main__':
    generate_md()
