import os
from io import BytesIO

from docx import Document
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
from langchain_anthropic import ChatAnthropic
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate

from prompts.generate_final_report_prompt import generate_final_report_prompt
from prompts.generate_initial_plan_prompt import generate_initial_plan_prompt
from prompts.generate_plan_prompt import generate_plan_prompt
from prompts.generate_report_format_prompt import generate_report_format_prompt
from prompts.report_format_prompt import format_prompt_dict

claude = ChatAnthropic(model='claude-3-5-sonnet-20240620', max_tokens=4096)

# 사용자의 질문에 대한 답변
generate_plan_response_schemas = [
    ResponseSchema(name="1", description="1번째 탐구 내용"),
    ResponseSchema(name="2", description="2번째 탐구 내용"),
    ResponseSchema(name="3", description="3번째 탐구 내용"),
    ResponseSchema(name="4", description="4번째 탐구 내용"),
    ResponseSchema(name="5", description="5번째 탐구 내용"),
    ResponseSchema(name="6", description="6번째 탐구 내용"),
    ResponseSchema(name="7", description="7번째 탐구 내용"),
    ResponseSchema(name="8", description="8번째 탐구 내용"),
]

generate_plan_format_response_schemas = [
    ResponseSchema(name="양식", description="탐구계획과 어울리는 양식의 번호 ex. 양식1"),
    ResponseSchema(name="이유", description="탐구계획과 어울리는 양식의 이유"),
]

generate_initial_plan_response_schemas = [
    ResponseSchema(name="문제상황", description="문제상황"),
    ResponseSchema(name="윤리적_고찰", description="윤리적 고찰"),
    ResponseSchema(name="결론", description="결론"),
]

# 응답 스키마를 기반으로 한 구조화된 출력 파서 초기화
# TODO: OutputParser 추가하지 않고, json 형태로 받아서 class로 변환하는 작업 필요
output_parser = StructuredOutputParser.from_response_schemas(generate_plan_response_schemas)

generate_plan_format_output_parser =   StructuredOutputParser.from_response_schemas(generate_plan_format_response_schemas)

generate_initial_plan_response_parser = StructuredOutputParser.from_response_schemas(generate_initial_plan_response_schemas)

# 출력 형식 지시사항을 파싱합니다.
generate_plan_format_instructions = output_parser.get_format_instructions()

generate_plan_prompt_template = PromptTemplate(
    template = generate_plan_prompt,
)

generate_format_prompt_template =  PromptTemplate(
    template = generate_report_format_prompt,
)

generate_initial_plan_prompt_template = PromptTemplate(
    template = generate_initial_plan_prompt,
)

generate_final_report_prompt_template = PromptTemplate(
    template = generate_final_report_prompt,
)

def generate_plans(ingredients: str):
    chain = generate_plan_prompt_template | claude | output_parser
    result = chain.invoke({"소재": ingredients, "format_instructions": generate_plan_format_instructions})
    print(result)
    return result

def manual(plans: str):
    chain = generate_format_prompt_template | claude | generate_plan_format_output_parser
    result = chain.invoke({"계획":plans, "format_instructions": generate_plan_format_output_parser.get_format_instructions()})
    print(result)
    return result
def generate_initial_plan(plans: str, guide_lines: str):
    chain = generate_initial_plan_prompt_template | claude | generate_initial_plan_response_parser
    result = chain.invoke({"계획":plans, "가이드라인":guide_lines, "format_instructions":generate_initial_plan_response_parser.get_format_instructions()})
    print(result)
    return result

def generate_final_plan(plans: str, problem: str, ethical_analysis: str, conclusion: str):
    chain = generate_final_report_prompt_template | claude | StrOutputParser()
    result = chain.invoke({"계획":plans, "문제상황":problem, "윤리적_고찰":ethical_analysis, "결론":conclusion})
    print(result)
    return result

def make_report(plans: str):
    format_result = manual(plans)
    format = format_prompt_dict[format_result['양식']]
    plan = generate_initial_plan(plans, format)
    return generate_final_plan(plans, plan['문제상황'], plan['윤리적_고찰'], plan['결론'])

def create_word_file(code):
    doc = Document()
    doc.add_heading('윤리 보고서', 0)  # 제목 추가
    doc.add_paragraph(code)  # 코드 내용을 워드 문서에 추가

    # BytesIO를 사용하여 메모리에 파일 저장
    file_stream = BytesIO()
    doc.save(file_stream)
    file_stream.seek(0)  # 파일 포인터를 처음으로 이동
    return file_stream