import base64
from enum import Enum
from typing import Iterator
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from config import SYSTEM_PROMPT
import json
import time


class Eligibility(str, Enum):
    ELIGIBLE = "Eligible"
    INELIGIBLE = "Ineligible"
    RESUBMIT = "Re-submit"


# class Response(BaseModel):
#     eligibility: Eligibility = Field(description="The status of eligibility One of: 'Eligible', 'Ineligible', 'Re-submit'.")
#     justification: str = Field(
#         description="A detailed, itemized, and numbered Markdown list of all findings that support the eligibility status. For 'Ineligible', list all violations. For 'Re-submit', list all points of uncertainty. For 'Eligible', state that no violations were found and note any movable items observed."
#     )
#     instruction: str = Field(
#         description="A clear, actionable, and numbered Markdown list of steps the homeowner must take. Each instruction must directly correspond to a finding in the 'justification' field. For 'Eligible' status, this provides a confirmation and reminders for movable items."
#     )

class Response(BaseModel):
    """
    Schema for the AI Inspector’s output.

    • eligibility  – Overall result.
    • justification – Numbered Markdown list explaining that result.
    • instruction  – Numbered Markdown list of next steps that map
                     one-to-one to the justification list.
    """

    eligibility: Eligibility = Field(
        ...,
        description=(
            "Overall inspection outcome.\n\n"
            " • 'Eligible'   – No violations observed in visible areas.\n"
            " • 'Ineligible' – ≥1 confirmed violation detected.\n"
            " • 'Re-submit'  – No clear violations, but ≥1 point of "
            "uncertainty prevents a confident decision."
        ),
    )

    justification: str = Field(
        ...,
        description=(
            "A *Markdown* numbered list detailing findings that back up "
            "the chosen eligibility value.\n\n"
            "Formatting rules:\n"
            " • Each item starts with '1.', '2.', …\n"
            " • For 'Eligible': a single sentence stating no violations "
            "   plus any movable combustibles noted.\n"
            " • For 'Ineligible': each item describes one specific, "
            "   observable violation.\n"
            " • For 'Re-submit': each item describes one uncertainty "
            "   that must be resolved (e.g., blurry photo, hidden area)."
        ),
        min_length=1,
    )

    instruction: str = Field(
        ...,
        description=(
            "A *Markdown* numbered list of homeowner actions that map "
            "directly, in order, to the items in 'justification'.\n\n"
            " • For 'Eligible': give a short confirmation plus reminders "
            "   to relocate any listed movable items during Red Flag "
            "   Warnings or extended absences.\n"
            " • For 'Ineligible': give a clear, actionable fix for each "
            "   violation.\n"
            " • For 'Re-submit': specify exactly what the homeowner must "
            "   provide or clarify (e.g., new close-up photo, daylight "
            "   image, measurement)."
        ),
        min_length=1,
    )

class LLMHandler:
    def __init__(self, model_name="gpt-5.2", temperature=0.3):
        self.llm = ChatOpenAI(model=model_name, temperature=temperature, streaming=True)
        self.parser = PydanticOutputParser(pydantic_object=Response)
        self.prompt = ChatPromptTemplate.from_messages(
            [
                ("system", SYSTEM_PROMPT),
                ("user", [
                    {
                        "type": "image_url",
                        "image_url": "data:image/jpeg;base64,{image_data}"
                    }
                ])
            ]
        ).partial(format_instructions=self.parser.get_format_instructions())
        self.chain = self.prompt | self.llm

    def process_image(self, image_path: str) -> Iterator[str]:
        with open(image_path, "rb") as img_file:
            image_data = base64.b64encode(img_file.read()).decode("utf-8")

        response = self.chain.invoke(
            {
                "image_data": image_data,
                "image_mime_type": "image/jpeg",
                "cache_type": "ephemeral",
            }
        )
        return response.text()

import gradio as gr

llm_handler = LLMHandler()


def process_and_display(image):
    if image is None:
        return "Please upload an image."

    image_path = image

    llm_output = json.loads(llm_handler.process_image(image_path))

    return llm_output["justification"], llm_output["instruction"]

demo = gr.Interface(
    fn=process_and_display,
    inputs=gr.Image(type="filepath", label="Upload Photo", height=500),
    outputs=[gr.Textbox(label="Inspection Result", info="The list of potential violations", lines=8),
             gr.Textbox(label="Homeowner Instruction", info="The instruction for how to fix the potential violations or what new evidence to provide." ,  lines=8)],
    title="Wildfire Prepared Home Eligibility Inspector",
    description="Upload a photo of all four sides of their home (one by one), showcasing the 0- to 5-foot noncombustible zone.",
    flagging_mode="never",
)

demo.launch()