from __future__ import annotations

import json
import os
import time
import urllib.request
from dataclasses import dataclass
from pathlib import Path

from playwright.sync_api import Page, expect, sync_playwright

BASE_URL = os.environ.get("KINK_AUDIT_BASE_URL", "http://127.0.0.1:8011")
ARTIFACT_DIR = Path("/tmp/kink_ui_audit")
ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
BONDAGE_ID = "fetlife_10"
SPANKING_ID = "fetlife_96"
ORAL_ID = "fetlife_71"


@dataclass
class Credentials:
    user_id: str
    private_token: str


def auth_headers(creds: Credentials) -> dict[str, str]:
    return {"x-private-token": creds.private_token}


def post_json(path: str, payload: dict, headers: dict[str, str] | None = None, timeout: int = 20) -> dict:
    request_headers = {"Content-Type": "application/json"}
    if headers:
        request_headers.update(headers)
    request = urllib.request.Request(
        f"{BASE_URL}{path}",
        data=json.dumps(payload).encode("utf-8"),
        headers=request_headers,
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def create_user() -> Credentials:
    payload = post_json("/users", {})
    return Credentials(user_id=payload["id"], private_token=payload["private_token"])


def save_play(creds: Credentials, kink_id: str, interest_state: str, directions: list[str]) -> dict:
    return post_json(
        f"/users/{creds.user_id}/plays",
        {"kink_id": kink_id, "interest_state": interest_state, "directions": directions},
        headers=auth_headers(creds),
    )


def link_partner(creds: Credentials, partner_id: str) -> dict:
    return post_json(
        f"/users/{creds.user_id}/partners",
        {"partner_id": partner_id},
        headers=auth_headers(creds),
    )


def create_partner_group(creds: Credentials, name: str, member_ids: list[str]) -> dict:
    return post_json(
        f"/users/{creds.user_id}/partner-groups",
        {"name": name, "member_ids": member_ids},
        headers=auth_headers(creds),
    )


def seed_audit_users() -> tuple[Credentials, Credentials, str]:
    owner = create_user()
    partner = create_user()
    save_play(owner, BONDAGE_ID, "love", ["to_me"])
    save_play(partner, BONDAGE_ID, "love", ["by_me"])
    save_play(owner, SPANKING_ID, "like", ["together"])
    save_play(partner, SPANKING_ID, "love", ["together"])
    save_play(owner, ORAL_ID, "love", ["together"])
    save_play(partner, ORAL_ID, "love", ["together"])
    link_partner(owner, partner.user_id)
    group = create_partner_group(owner, "UI Audit Duo", [partner.user_id])
    return owner, partner, group["id"]


def capture_locator(locator, label: str, path: Path) -> dict:
    if locator.count() == 0:
        return {"label": label, "present": False}
    target = locator.first
    box = target.bounding_box()
    if not box:
        return {"label": label, "present": True, "box": None}
    target.screenshot(path=str(path))
    return {
        "label": label,
        "present": True,
        "box": {
            "x": round(box["x"]),
            "y": round(box["y"]),
            "width": round(box["width"]),
            "height": round(box["height"]),
        },
        "path": str(path),
    }


def capture_layout_state(page: Page) -> dict:
    return page.evaluate(
        """() => ({
          viewport: { width: window.innerWidth, height: window.innerHeight },
          bodyScrollHeight: document.body.scrollHeight,
          headings: [...document.querySelectorAll('h1, h2, h3')].map((el) => ({
            text: (el.textContent || '').trim(),
            top: Math.round(el.getBoundingClientRect().top),
          })).slice(0, 20),
          statusText: [...document.querySelectorAll('.status')].map((el) => (el.textContent || '').trim()).filter(Boolean),
          authText: document.querySelector('.auth-strip')?.innerText || '',
          selectedDetailTitle: document.querySelector('.detail-title')?.textContent || ''
        })"""
    )


def round_box(box: dict | None) -> dict | None:
    if not box:
        return None
    return {
        "x": round(box["x"]),
        "y": round(box["y"]),
        "width": round(box["width"]),
        "height": round(box["height"]),
    }


def first_box(locator) -> dict | None:
    if locator.count() == 0:
        return None
    return round_box(locator.first.bounding_box())


def assert_close(label: str, left: int, right: int, tolerance: int = 12) -> None:
    if abs(left - right) > tolerance:
        raise AssertionError(f"{label} misaligned by {abs(left - right)}px ({left} vs {right})")


def capture_alignment_metrics(page: Page) -> dict:
    add_panel = page.locator('.panel').filter(has=page.get_by_role('heading', name='Add Plays')).first
    for_you_panel = page.locator('.panel').filter(has=page.get_by_role('heading', name='For You', exact=True)).first
    shared_panel = page.locator('.panel').filter(has=page.get_by_role('heading', name='Shared Play List')).first
    detail_pane = page.locator('.side-stack').first
    board_columns = page.locator('.board-grid .board-column')

    add_box = first_box(add_panel)
    for_you_box = first_box(for_you_panel)
    shared_box = first_box(shared_panel)
    detail_box = first_box(detail_pane)
    board_boxes = [round_box(board_columns.nth(i).bounding_box()) for i in range(board_columns.count())]
    board_boxes = [box for box in board_boxes if box]

    add_heading = first_box(add_panel.locator('h2, h3').first) if add_panel.count() else None
    for_you_heading = first_box(for_you_panel.locator('h2, h3').first) if for_you_panel.count() else None
    shared_heading = first_box(shared_panel.locator('h2, h3').first) if shared_panel.count() else None

    return {
        "panels": {
            "add": add_box,
            "for_you": for_you_box,
            "shared": shared_box,
            "detail": detail_box,
        },
        "headings": {
            "add": add_heading,
            "for_you": for_you_heading,
            "shared": shared_heading,
        },
        "board_columns": board_boxes,
    }


def assert_shell_alignment(metrics: dict, *, mobile: bool) -> None:
    panels = metrics["panels"]
    headings = metrics["headings"]
    board_columns = metrics["board_columns"]

    required_panels = [panels["add"], panels["for_you"], panels["shared"]]
    if any(panel is None for panel in required_panels):
        raise AssertionError("missing panel box for alignment audit")
    if any(heading is None for heading in [headings["add"], headings["for_you"], headings["shared"]]):
        raise AssertionError("missing heading box for alignment audit")

    assert_close("add vs for_you left edge", panels["add"]["x"], panels["for_you"]["x"])
    assert_close("for_you vs shared left edge", panels["for_you"]["x"], panels["shared"]["x"])
    assert_close("add vs for_you heading inset", headings["add"]["x"] - panels["add"]["x"], headings["for_you"]["x"] - panels["for_you"]["x"])
    assert_close("for_you vs shared heading inset", headings["for_you"]["x"] - panels["for_you"]["x"], headings["shared"]["x"] - panels["shared"]["x"])

    if not mobile:
        if len(board_columns) != 3:
            raise AssertionError(f"expected 3 board columns, got {len(board_columns)}")
        widths = [box["width"] for box in board_columns]
        heights = [box["height"] for box in board_columns]
        left_edges = [box["x"] for box in board_columns]
        if max(widths) - min(widths) > 8:
            raise AssertionError(f"board column widths diverge too much: {widths}")
        if max(heights) - min(heights) > 20:
            raise AssertionError(f"board column heights diverge too much: {heights}")
        gaps = [left_edges[i + 1] - left_edges[i] for i in range(len(left_edges) - 1)]
        if gaps and max(gaps) - min(gaps) > 12:
            raise AssertionError(f"board column spacing inconsistent: {left_edges}")


def login_storage_script(creds: Credentials) -> str:
    payload = json.dumps({"userId": creds.user_id, "token": creds.private_token})
    return f"localStorage.setItem('kink_auth', JSON.stringify({payload}));"


def run_scenario(name: str, viewport: tuple[int, int], init_script: str, action) -> dict:
    console_errors: list[str] = []
    start = time.perf_counter()
    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        page = browser.new_page(viewport={"width": viewport[0], "height": viewport[1]})
        page.add_init_script(init_script)
        page.on("console", lambda msg: console_errors.append(f"{msg.type}: {msg.text}") if msg.type == "error" else None)
        page.goto(BASE_URL, wait_until="domcontentloaded")
        page.wait_for_timeout(1200)
        result = {
            "name": name,
            "viewport": {"width": viewport[0], "height": viewport[1]},
            "console_errors": console_errors,
            "checkpoints": [],
        }
        action(page, result)
        full_path = ARTIFACT_DIR / f"{name}.png"
        page.screenshot(path=str(full_path), full_page=True)
        result["full_screenshot"] = str(full_path)
        result["elapsed_ms"] = round((time.perf_counter() - start) * 1000, 1)
        browser.close()
        return result


def action_authed_shell(page: Page, result: dict) -> None:
    expect(page.locator('.auth-strip')).to_be_visible(timeout=15000)
    expect(page.get_by_role('heading', name='For You', exact=True)).to_be_visible(timeout=15000)
    state = capture_layout_state(page)
    alignment = capture_alignment_metrics(page)
    assert_shell_alignment(alignment, mobile=page.viewport_size['width'] <= 480)
    result['checkpoints'].append({"label": "shell", "state": state, "alignment": alignment})
    result['clips'] = [
        capture_locator(page.locator('.panel').filter(has=page.get_by_role('heading', name='For You', exact=True)), 'for_you_panel', ARTIFACT_DIR / f"{result['name']}_for_you.png"),
        capture_locator(page.locator('.board-grid'), 'play_list_board', ARTIFACT_DIR / f"{result['name']}_play_list.png"),
        capture_locator(page.locator('.panel').filter(has=page.get_by_role('heading', name='Shared Play List')), 'shared_panel', ARTIFACT_DIR / f"{result['name']}_shared_panel.png"),
    ]


def action_detail_open(page: Page, result: dict) -> None:
    expect(page.locator('.auth-strip')).to_be_visible(timeout=15000)
    page.get_by_placeholder('Search plays').fill('bondage')
    chip = page.locator('.panel').filter(has=page.get_by_text('Add Plays')).locator('.play-chip').filter(has=page.get_by_role('button', name='Bondage')).first
    expect(chip).to_be_visible(timeout=15000)
    chip.get_by_role('button', name='Bondage').click()
    expect(page.locator('.detail-title')).to_have_text('Bondage', timeout=15000)
    result['checkpoints'].append({"label": "detail_open", "state": capture_layout_state(page), "alignment": capture_alignment_metrics(page)})
    result['clips'] = [
        capture_locator(page.locator('.side-stack'), 'detail_pane', ARTIFACT_DIR / f"{result['name']}_detail.png"),
        capture_locator(page.locator('.panel').filter(has=page.get_by_role('heading', name='For You', exact=True)), 'for_you_panel', ARTIFACT_DIR / f"{result['name']}_for_you.png"),
    ]


def action_shared_group(page: Page, result: dict) -> None:
    expect(page.locator('.auth-strip')).to_be_visible(timeout=15000)
    shared_panel = page.locator('.panel').filter(has=page.get_by_role('heading', name='Shared Play List'))
    shared_panel.get_by_role('button', name='UI Audit Duo (2)').click()
    expect(page.locator('.shared-grid')).to_be_visible(timeout=15000)
    result['checkpoints'].append({"label": "shared_group", "state": capture_layout_state(page)})
    result['clips'] = [
        capture_locator(page.locator('.shared-grid'), 'shared_grid', ARTIFACT_DIR / f"{result['name']}_shared.png"),
        capture_locator(shared_panel.locator('.chips').first, 'group_picker', ARTIFACT_DIR / f"{result['name']}_group_picker.png"),
    ]


def main() -> int:
    owner, _partner, _group_id = seed_audit_users()
    scenarios = [
        run_scenario('desktop_authed_shell', (1440, 1000), login_storage_script(owner), action_authed_shell),
        run_scenario('desktop_detail_open', (1440, 1000), login_storage_script(owner), action_detail_open),
        run_scenario('desktop_shared_group', (1440, 1000), login_storage_script(owner), action_shared_group),
        run_scenario('mobile_authed_shell', (390, 844), login_storage_script(owner), action_authed_shell),
        run_scenario('mobile_detail_open', (390, 844), login_storage_script(owner), action_detail_open),
    ]
    out = {
        'generated_at': time.time(),
        'base_url': BASE_URL,
        'artifact_dir': str(ARTIFACT_DIR),
        'scenarios': scenarios,
    }
    out_path = ARTIFACT_DIR / 'audit.json'
    out_path.write_text(json.dumps(out, indent=2))
    print(json.dumps({
        'artifact_dir': str(ARTIFACT_DIR),
        'audit_json': str(out_path),
        'scenarios': [
            {
                'name': item['name'],
                'elapsed_ms': item['elapsed_ms'],
                'console_errors': item['console_errors'],
                'clip_count': len(item.get('clips', [])),
            }
            for item in scenarios
        ],
    }, indent=2))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())