import os, json, gradio as gr
from PIL import Image
import pytesseract

# --- import your logic ---
from checker import (
    evaluate_section,
    fair_housing_flags,
    COMPANY_NAME_DEFAULT,
    COMPANY_PHONES_DEFAULT,
    DISCLAIMER_DEFAULT,
    contains_disclaimer,
    count_name_instances,
    count_phone_instances,
)

def run_check(image, ptxt, social, agent_name, agent_phone,
              company_name, company_phones_json, disclaimer):
    # OCR
    itxt = ""
    ocr_err = None
    if image is not None:
        try:
            itxt = pytesseract.image_to_string(image)
        except Exception as e:
            ocr_err = str(e)

    # Compose combined content
    content = "\n\n".join([x for x in [itxt, ptxt, f"Social={social}"] if x])

    # Fair-housing flags on combined content
    fh_flags = fair_housing_flags(content)
    fair_housing_block = {"compliant": len(fh_flags) == 0, "Flags": fh_flags}

    # Parse office phones
    try:
        company_phones = json.loads(company_phones_json)
        if isinstance(company_phones, str):
            company_phones = [company_phones]
    except Exception:
        company_phones = COMPANY_PHONES_DEFAULT

    # Social disclaimer toggle (same behavior as your prototype)
    require_disclaimer_on_social = os.getenv("REQUIRE_DISCLAIMER_ON_SOCIAL", "1") == "1"

    def eval_section(text):
        flags = []

        company_name_count = count_name_instances(text, company_name)
        agent_name_count   = count_name_instances(text, agent_name)
        office_phone_count = count_phone_instances(text, company_phones)
        agent_phone_count  = count_phone_instances(text, [agent_phone] if agent_phone else [])

        name_equal  = (company_name_count == agent_name_count)
        phone_equal = (office_phone_count == agent_phone_count)

        disclaimer_ok = True
        if social and require_disclaimer_on_social:
            disclaimer_ok = contains_disclaimer(text, disclaimer)
            if not disclaimer_ok:
                flags.append("Missing disclaimer on social content")

        if not name_equal:
            flags.append(f"Name imbalance: company={company_name_count} vs agent={agent_name_count}")
        if not phone_equal:
            flags.append(f"Phone imbalance: office={office_phone_count} vs agent={agent_phone_count}")

        compliant = name_equal and phone_equal and disclaimer_ok
        return {"compliant": compliant, "Flags": flags}

    img_block  = eval_section(itxt)
    if ocr_err:
        img_block["Flags"].append(f"OCR error: {ocr_err}")

    ptxt_block = eval_section(ptxt or "")

    # final payload in your exact shape
    payload = {
        "Fair_Housing": fair_housing_block,
        "img": img_block,
        "Ptxt": ptxt_block
    }
    return json.dumps(payload, indent=2)

with gr.Blocks(title="Image + Text Compliance Check") as demo:
    gr.Markdown("# Image + Text Compliance Check")

    with gr.Row():
        image = gr.Image(type="pil", label="Upload image (optional)")
        ptxt  = gr.Textbox(lines=8, label="Post Text (Ptxt)")

    with gr.Row():
        social      = gr.Checkbox(label="Social", value=False)
        agent_name  = gr.Textbox(label="Agent Name", placeholder="e.g., Jane Doe")
        agent_phone = gr.Textbox(label="Agent Phone (digits or formatted)")

    with gr.Accordion("Advanced", open=False):
        company_name        = gr.Textbox(label="Company Name", value=COMPANY_NAME_DEFAULT)
        company_phones_json = gr.Textbox(label="Company Phones (JSON list)", value=json.dumps(COMPANY_PHONES_DEFAULT))
        disclaimer          = gr.Textbox(label="Disclaimer", value=DISCLAIMER_DEFAULT)

    run_btn = gr.Button("Run Compliance Check")
    out     = gr.Code(label="Result JSON", language="json")

    run_btn.click(
        fn=run_check,
        inputs=[image, ptxt, social, agent_name, agent_phone, company_name, company_phones_json, disclaimer],
        outputs=[out],
    )

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