from ocr.api.prompts import ocr_prompts
from ocr.api.report.model import ReportModel
from ocr.core.wrappers import openai_wrapper


@openai_wrapper()
async def generate_report(request_content: str):
    messages = [
        {
            "role": "system",
            "content": ocr_prompts.report.generate_report
        },
        {
            "role": "user",
            "content": request_content
        }
    ]
    return messages


@openai_wrapper()
async def generate_changes(content: str, previous_report: str):
    messages = [
        {
            "role": "system",
            "content": ocr_prompts.report.generate_changes
            .replace("{previous_report}", previous_report)
        },
        {
            "role": "user",
            "content": content
        }
    ]
    return messages


@openai_wrapper()
async def generate_agent_response(messages: list[dict], report: ReportModel):
    messages = [
        {
            "role": "system",
            "content": ocr_prompts.message.generate_agent_response
            .replace("{reports}", report.report)
            .replace("{changes}", report.changes or 'There is no changes.')
        },
        *messages
    ]
    return messages


@openai_wrapper(is_json=True, temperature=0.6, return_='result')
async def generate_consult_note(text: str, changes: str, type_: str):
    prompt_map = {
        "chief": ocr_prompts.consult.generate_chief,
        "hpi": ocr_prompts.consult.generate_hpi,
        "social": ocr_prompts.consult.generate_social,
        "surgical": ocr_prompts.consult.generate_surgical,
        "family": ocr_prompts.consult.generate_family,
        "medications": ocr_prompts.consult.generate_medications,
        "assessment": ocr_prompts.consult.generate_assessment,
        "plan": ocr_prompts.consult.generate_plan,
    }
    user_content = f'Medical information:\n```\n{text}\n```'
    if changes:
        user_content += f'\nChanges:\n```\n{changes}\n```'
    messages = [
        {
            "role": "system",
            "content": prompt_map[type_]
        },
        {
            "role": "user",
            "content": user_content
        }
    ]
    return messages