import gradio as gr
import os
import requests
import json
import logging
from openai import OpenAI
import json
import textwrap
import requests
import xml.etree.ElementTree as ET
from lxml import etree

openAiClient = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY")
)
perplexityClient = OpenAI(
    api_key=os.getenv("PERPLEXITY_API_KEY"),
    base_url="https://api.perplexity.ai"
)

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

PERPLEXITY_API_URL = "https://api.perplexity.ai/chat/completions"
PERPLEXITY_API_KEY = os.getenv("PERPLEXITY_API_KEY")
with open('tax_rules.txt', 'r', encoding='utf-8') as file:
    tax_rules = file.read()

ANSWERS = {}

ASSISTANT_PROMPT = textwrap.dedent(f"""
Jestem asystentem do wypelnienia deklaracji PCC. Uzytkownik bedzie zadawac rozne pytania, 
na ktore odpowiedz zgodnie z prawem. Odpowiedzi udzielaj stosujac nastepujace zasady:

{tax_rules}

Odpowiedzi przeslij w formacie JSON bez zadnego innego tekstu albo formatowania w nastepujacym formacie:
{{
    "answer": "Odpowiedz na pytanie uzytkownika",
    "user_answer": "Odpowiedz uzytkownika na pytanie",
    "ready_for_next_question": "Czy uzytkownik jest gotowy na nastepne pytanie? Powinno być true tylko jesli juz wczesniej wyrazil gotowosc na wypelnienie deklaracji (true/false)",
    "ready_for_questionnare": "Jesli uzytkownik wyrazil chec na wypelnienie deklaracji, to powinno byc ustawione na true i tylko gdy uzytkownik wyraznie powie ze jest gotowy na wypelnienie deklaracji. (true/false)",
    "is_answering_questions": "Czy uzytkownik jest w trakcie podawania odpowiedzi na pytania do formularza? To pole powinno być ustawione na true tylko gdy uzytkownik jest w trakcie wypelniania deklaracji. (true/false)",
    "answer_correct": "Czy odpowiedz jest poprawna? Jesli uzytkownik poda odpowiedz, ktora jest niepoprawna, to powinno byc ustawione na false. (true/false)",
    "case_purpose": "Opis z czego uzytkownik chce sie rozliczyc aby wybrac poprawny formularz."
}}

Odpowiedz powinna zawierac jedynie poprawny JSON, a odpowiednie pola w JSON'ie w jezyku polskim. Asystent powinien na początku zadać pytanie o to z czego uzytkownik chce sie rozliczyc aby wybrac poprawny formularz.
Powinien równiez zapytać czy uzytkownik jest osoba fizyczna czy firma. W przypadku gdy uzytkownik jest zwolniony z podatku, poinformuj go o tym.
""")

def remove_markdown(text):
    text = text.replace('```json', '')
    text = text.replace('```', '')
 
    return text.strip()


def generate_xml_from_xsd(xsd_url, answers):
    response = requests.get(xsd_url)
    if response.status_code != 200:
        raise Exception(f"Failed to fetch XSD file from {xsd_url}")
    
    xsd_root = etree.fromstring(response.content)    
    schema = etree.XMLSchema(xsd_root)    
    parser = etree.XMLParser(schema=schema)
    root = ET.Element(xsd_root.find(".//xs:element", namespaces={"xs": "http://www.w3.org/2001/XMLSchema"}).get("name"))
    
    for field, value in answers.items():
        element = ET.SubElement(root, field)
        element.text = str(value)
    
    try:
        schema.assertValid(root)
    except etree.DocumentInvalid as e:
        print(f"Generated XML is not valid according to the schema: {e}")
    
    xml_string = ET.tostring(root, encoding="unicode", method="xml")
    
    return xml_string

def generate_pcc3_xml():
    global pcc3_structure, answers
    if pcc3_structure and 'xsd_file' in pcc3_structure:
        try:
            xml_string = generate_xml_from_xsd(pcc3_structure['xsd_file'], answers)
            return xml_string
        except Exception as e:
            logger.error(f"Error generating XML: {e}")
            return None
    else:
        logger.error("PCC3 structure or XSD file URL not available")
        return None


def fetch_pcc3_info(case_purpose):
    logger.info("Fetching PCC information to formulate questions")
    system_message = textwrap.dedent(f"""
    Jesteś asystentem AI, którego zadaniem jest pomoc użytkownikom w wypełnianiu formularza podatkowego PCC w Polsce.
    Przeanalizuj dostarczone informacje o formularzach PCC i stwórz ustrukturyzowany plan zbierania niezbędnych informacji od użytkowników. 
    Kazde pytanie powinno pytac tylko o jedna informacje. Uzytkownik powie z czego dokladnie chce sie rozliczyc aby wybrac odpowiedni formularz.
    Twoja odpowiedź powinna być w formacie JSON o następującej strukturze. Zastosuj nastepujace zasady co do wyboru wymaganego formularza:

    {tax_rules}

    {{
        "form_description": "A brief description of the PCC form",
        "xsd_file": "Direct link to the concrete XSD file that was used to create the questions. This link should not point to anything else than XSD file. ",
        "xsd_file_content": "Content of the XSD file. This should be a string that contains the content of the XSD file used to create the questions.",
        "form_id": "ID of the form that should be used to fill in the answers",
        "required_fields": [
            {{
                "field_name": "Name of the field in the xsd file",
                "description": "Description of what information is needed",
                "question": "A question to ask the user to gather this information"
            }},        
            ...
        ]
    }}
    Odpowiedz tylko JSON'em i niczym więcej. Nie stosuj zadnego formatowania. 
    Nie stosuj zadnego wstępu. Wybierz plik xsd ktory najbardziej pasuje do podanego przez uzytkownika zadania i uzyj go do stworzenia pytan do uzytkownika. 
    Pytania z formularza powinny miec takie same nazwy pola jak pola w pliku xsd. 
    Uzyj nastepujacej strony w celu znalezienia odpowiedniego formularza do wypelnienia: https://www.podatki.gov.pl/pcc-sd/e-deklaracje-pcc-sd/formularze-pcc/
    """)

    try:
        response = perplexityClient.chat.completions.create(
            model="llama-3.1-sonar-huge-128k-online",
            max_tokens=10000,
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": case_purpose}
            ],
        )
        
        return json.loads(remove_markdown(response.choices[0].message.content))  # Parse JSON here
    except requests.RequestException as e:
        logger.error(f"Error calling Perplexity API: {e}")
        return {"error": str(e)}
    except json.JSONDecodeError as e:
        logger.error(f"Error decoding JSON from Perplexity API response: {e}")
        return {"error": "Invalid JSON response from API"}

pcc3_structure = None
answers = {}
current_question_index = 0

def fetch_next_question():
    global current_question_index
    if current_question_index >= len(pcc3_structure["required_fields"]):
        return None
    question = pcc3_structure["required_fields"][current_question_index]["question"]
    current_question_index += 1
    return question

def fetch_answer(input_string):
    last_brace_index = input_string.rfind('}')
    
    if last_brace_index == -1:
        return input_string.strip()
    
    try:
        json_part = input_string[:last_brace_index + 1]
        json_data = json.loads(json_part)
        
        if "answer" in json_data:
            return json_data["answer"]
    except json.JSONDecodeError:
        pass
    
    remaining_text = input_string[last_brace_index + 1:].strip()
    return remaining_text if remaining_text else input_string.strip()

def gpt4_omni_assistant(message, history):
    global pcc3_structure
    global answers
    if history is None:
        history = []

    messages = [{"role": "system", "content": ASSISTANT_PROMPT}]
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})

    messages.append({"role": "user", "content": message})
    response = openAiClient.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        max_tokens=1000
    )
    assistant_msg = json.loads(remove_markdown(response.choices[0].message.content))
    if assistant_msg["ready_for_questionnare"] == True and pcc3_structure == None:
        logger.info("Fetching PCC-3 structure")
        pcc3_structure = fetch_pcc3_info(assistant_msg["case_purpose"])
        messages.append({
            "role": "system",
            "content": "Zadaj uzytkownikowi pytanie: " + fetch_next_question()
        })
        return fetch_answer(json.loads(remove_markdown(openAiClient.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            max_tokens=1000
        ).choices[0].message.content))["answer"])
    
    if assistant_msg["answer_correct"] == True and assistant_msg["is_answering_questions"] == True:
        logger.info("User answered correctly: " + assistant_msg["user_answer"])
        ANSWERS[pcc3_structure["required_fields"][current_question_index - 1]["field_name"]] = assistant_msg["user_answer"]
        answers[pcc3_structure["required_fields"][current_question_index - 1]["field_name"]] = assistant_msg["user_answer"]

    if assistant_msg["ready_for_next_question"] == True and assistant_msg["is_answering_questions"] == True:
        next_question = fetch_next_question()
        if next_question is None:
            return "Dziekujemy za wypelnienie deklaracji. Jesli masz dodatkowe pytania, to pytaj."
        return next_question
    return fetch_answer(assistant_msg["answer"])

demo = gr.ChatInterface(
    fn=gpt4_omni_assistant,
    title="Asystent Podatkowy",
    description="Odpowiem na wszystkie pytania dotyczące podatków i pomogę Ci w wypelnieniu deklaracji!",
    theme=gr.themes.Default(),
    css=".gradio-container {height: 100vh !important;}",
    chatbot=gr.Chatbot(placeholder=f"Dzień dobry, jestem tutaj aby pomóc Ci w wypełnieniu deklaracji PCC-3."),    
)

if __name__ == "__main__":
    demo.launch()