import os
import re
import socket
import ipaddress
from functools import lru_cache
from typing import Optional
from urllib.parse import urlparse

import gradio as gr
import requests
import torch
import uvicorn
from bs4 import BeautifulSoup
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, ConfigDict, Field, model_validator
from transformers import AutoModelForCausalLM, AutoTokenizer


APP_NAME = "Truth"
SPACE_URL = "https://adedoyinjames-truth.hf.space"
AUTHOR = "Adedoyin Ifeoluwa JAmes"
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-0.5B-Instruct")
MAX_INPUT_CHARS = int(os.getenv("MAX_INPUT_CHARS", "12000"))
MAX_FETCH_BYTES = int(os.getenv("MAX_FETCH_BYTES", "1500000"))
DEFAULT_MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "320"))


class SummarizeRequest(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    text: Optional[str] = Field(
        default=None,
        alias="policy_text",
        description="Raw privacy policy, terms, or legal text to summarize.",
    )
    url: Optional[str] = Field(
        default=None,
        alias="policy_url",
        description="Public HTTP(S) URL to fetch and summarize.",
    )
    source_hint: Optional[str] = Field(
        default=None,
        description="Optional app or website name, such as Chrome-detected signup page.",
    )
    max_new_tokens: int = Field(default=DEFAULT_MAX_NEW_TOKENS, ge=80, le=700)

    @model_validator(mode="after")
    def require_text_or_url(self):
        if not (self.text and self.text.strip()) and not (self.url and self.url.strip()):
            raise ValueError("Provide either text or url.")
        return self


class SummarizeResponse(BaseModel):
    app: str
    model: str
    source: str
    summary: str
    bullets: list[str]


def _is_public_http_url(url: str) -> bool:
    parsed = urlparse(url)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        return False

    try:
        addresses = socket.getaddrinfo(parsed.hostname, None)
    except socket.gaierror:
        return False

    for address in addresses:
        ip = ipaddress.ip_address(address[4][0])
        if (
            ip.is_private
            or ip.is_loopback
            or ip.is_link_local
            or ip.is_multicast
            or ip.is_reserved
            or ip.is_unspecified
        ):
            return False
    return True


def fetch_policy_text(url: str) -> str:
    if not _is_public_http_url(url):
        raise HTTPException(status_code=400, detail="Only public HTTP(S) URLs are allowed.")

    headers = {
        "User-Agent": f"{APP_NAME}/1.0 policy summarizer ({SPACE_URL})",
        "Accept": "text/html,text/plain,application/xhtml+xml",
    }

    try:
        with requests.get(url, headers=headers, timeout=15, stream=True) as response:
            response.raise_for_status()
            chunks: list[bytes] = []
            total = 0
            for chunk in response.iter_content(chunk_size=8192):
                if not chunk:
                    continue
                total += len(chunk)
                if total > MAX_FETCH_BYTES:
                    break
                chunks.append(chunk)
            raw = b"".join(chunks)
            content_type = response.headers.get("content-type", "")
    except requests.RequestException as exc:
        raise HTTPException(status_code=502, detail=f"Could not fetch URL: {exc}") from exc

    decoded = raw.decode(response.encoding or "utf-8", errors="ignore")
    if "html" in content_type or "<html" in decoded[:500].lower():
        soup = BeautifulSoup(decoded, "html.parser")
        for tag in soup(["script", "style", "noscript", "svg"]):
            tag.decompose()
        text = soup.get_text(" ", strip=True)
    else:
        text = decoded

    text = re.sub(r"\s+", " ", text).strip()
    if not text:
        raise HTTPException(status_code=422, detail="Fetched URL did not contain readable text.")
    return text


@lru_cache(maxsize=1)
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype="auto",
        device_map="auto",
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )
    model.eval()
    return tokenizer, model


def build_prompt(policy_text: str, source_hint: Optional[str]) -> str:
    source_line = f"Source hint: {source_hint.strip()}" if source_hint else "Source hint: Not provided"
    trimmed = policy_text[:MAX_INPUT_CHARS]
    messages = [
        {
            "role": "system",
            "content": (
                "You are Truth, a privacy policy and terms-of-service summarizer. "
                "Use only the provided text. Explain risks plainly for a mobile user. "
                "Return exactly four bullets. Do not include legal advice."
            ),
        },
        {
            "role": "user",
            "content": (
                f"{source_line}\n\n"
                "Summarize this policy in exactly four short bullets:\n"
                "1. What data is collected.\n"
                "2. Who data is shared with or sold to.\n"
                "3. What controls, rights, or opt-outs the user has.\n"
                "4. The biggest privacy concern in plain English.\n\n"
                f"Policy text:\n{trimmed}"
            ),
        },
    ]

    tokenizer, _ = load_model()
    if hasattr(tokenizer, "apply_chat_template"):
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    return "\n\n".join(f"{message['role'].upper()}: {message['content']}" for message in messages)


def parse_bullets(summary: str) -> list[str]:
    lines = [line.strip(" -*\t") for line in summary.splitlines()]
    bullets = []
    for line in lines:
        cleaned = re.sub(r"^\d+[\).\s-]+", "", line).strip()
        if cleaned:
            bullets.append(cleaned)
    return bullets[:4]


def generate_summary(
    text: Optional[str],
    url: Optional[str],
    source_hint: Optional[str],
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
) -> SummarizeResponse:
    source = "text"
    policy_text = (text or "").strip()
    if not policy_text and url:
        source = url.strip()
        policy_text = fetch_policy_text(source)

    if len(policy_text) < 80:
        raise HTTPException(status_code=422, detail="Policy text is too short to summarize reliably.")

    tokenizer, model = load_model()
    prompt = build_prompt(policy_text, source_hint)
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096)
    inputs = {key: value.to(model.device) for key, value in inputs.items()}

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    new_tokens = output_ids[0][inputs["input_ids"].shape[-1] :]
    summary = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    bullets = parse_bullets(summary)

    if len(bullets) < 4:
        bullets = (bullets + [summary])[:4]

    return SummarizeResponse(
        app=APP_NAME,
        model=MODEL_ID,
        source=source,
        summary="\n".join(f"- {bullet}" for bullet in bullets),
        bullets=bullets,
    )


api = FastAPI(
    title=f"{APP_NAME} API",
    description="Privacy policy summarizer for the Truth Android accessibility companion.",
    version="1.0.0",
)

api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)


@api.get("/health")
def health():
    return {
        "ok": True,
        "app": APP_NAME,
        "author": AUTHOR,
        "model": MODEL_ID,
    }


@api.post("/summarize", response_model=SummarizeResponse)
@api.post("/v1/summarize", response_model=SummarizeResponse)
def summarize(request: SummarizeRequest):
    return generate_summary(
        text=request.text,
        url=request.url,
        source_hint=request.source_hint,
        max_new_tokens=request.max_new_tokens,
    )


def summarize_for_ui(text: str, url: str, source_hint: str):
    try:
        response = generate_summary(
            text=text,
            url=url,
            source_hint=source_hint,
            max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
        )
        return response.summary
    except HTTPException as exc:
        return f"Error: {exc.detail}"
    except Exception as exc:
        return f"Error: {exc}"


demo = gr.Interface(
    fn=summarize_for_ui,
    inputs=[
        gr.Textbox(label="Policy text", lines=10, placeholder="Paste privacy policy or terms text here..."),
        gr.Textbox(label="Policy URL", placeholder="https://example.com/privacy"),
        gr.Textbox(label="Source hint", placeholder="Example: Chrome signup page"),
    ],
    outputs=gr.Textbox(label="Truth summary", lines=8),
    title="Truth",
    description=(
        "Truth summarizes privacy policies into four plain-English bulletins for an Android "
        "accessibility companion."
    ),
    flagging_mode="never",
)

app = gr.mount_gradio_app(api, demo, path="/")


if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)