# ─────────────────────────────────────────────────────────────────
# BizPortal 2.0 — Python Playwright sidecar
#
# Tiny FastAPI service that drives https://www.bizportal.gov.za via
# a real headless Chromium. The Next.js side talks to this over HTTP
# (no LLM agent, no file queue, no polling). All UI selectors live
# here so the existing /bizportal page is untouched.
#
# Run:
#   cd scripts
#   python3 -m venv .venv && source .venv/bin/activate
#   pip install -r requirements.txt
#   playwright install chromium
#   cp .env.example .env   # fill in credentials
#   uvicorn bizportal_service:app --host 127.0.0.1 --port 7001
#
# Smoke test:
#   curl http://127.0.0.1:7001/health
#   curl -X POST http://127.0.0.1:7001/search \
#     -H 'Content-Type: application/json' \
#     -d '{"type":"ck","query":"2022/808158/08"}'
#
# Endpoints:
#   GET  /health
#   POST /search        { type: "id"|"ck", query: str }
#   POST /reserve       { names: [str,str,str,str], dry_run: bool, balance_min: float=50 }
# ─────────────────────────────────────────────────────────────────

import os
import re
import json
import asyncio
import logging
from pathlib import Path
from typing import List, Optional, Dict, Any

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from playwright.async_api import async_playwright, Browser, BrowserContext, Page, TimeoutError as PWTimeout

load_dotenv()

# ─────────────────────────────────────────────────────────────────
# Config — credentials, endpoints, selectors. Tweak here, not in code.
# ─────────────────────────────────────────────────────────────────

BIZPORTAL_BASE = "https://www.bizportal.gov.za"
LOGIN_URL = f"{BIZPORTAL_BASE}/login.aspx"
PROFILE_URL = f"{BIZPORTAL_BASE}/user_profile.aspx"
NAME_RESERVE_URL = f"{BIZPORTAL_BASE}/name_reserve.aspx"
BIZPROFILE_URL = f"{BIZPORTAL_BASE}/bizprofile.aspx"

# Login fields — verified.
LOGIN_ID = "cntMain_txtIDNo"
LOGIN_PASSWORD = "cntMain_txtPassword"
LOGIN_BUTTON = "cntMain_btnLogin"

# Name reservation fields — verified.
NAME_INPUTS = ["cntMain_txtPropName1", "cntMain_txtPropName2", "cntMain_txtPropName3", "cntMain_txtPropName4"]
NAME_CONTINUE_1 = "cntMain_btnContinuePropName"
NAME_CONTINUE_2 = "cntMain_btnContinueLodge"

# BizProfile search — verified. ONE page: dropdown picks the search
# type, single input + button. Verified 2026-08-17 against the live
# site by probing with Playwright.
BIZPROFILE_DROPDOWN = "cntMain_drpSearchOptions"
BIZPROFILE_INPUT = "cntMain_txtSearchCIPC"
BIZPROFILE_BUTTON = "cntMain_btnSearch"
BIZPROFILE_DROPDOWN_OPTIONS = {
    "id": "IDNo",         # Director ID No.
    "ck": "EntNo",        # Enterprise No.
    "name": "EntName",    # Enterprise Name (unused for now)
}

# Company details labels on bizprofile.aspx (CK search result page).
BIZPROFILE_LBL_ENT_NO = "cntMain_lblEntNo"
BIZPROFILE_LBL_ENT_NAME = "cntMain_lblEntName"
BIZPROFILE_LBL_ENT_TYPE = "cntMain_lblEntType"
BIZPROFILE_LBL_ENT_STATUS = "cntMain_lblEntStatus"

# ID search returns a GridView with this id.
ID_SEARCH_GRID = "cntMain_gdvEnterprises"

# Maximum number of companies to click-through for director lookup
# during an ID search. The existing UI can handle more, but each
# click is a postback round-trip — cap it to keep searches fast.
MAX_ID_SEARCH_DIRECTOR_FETCH = 10

# Session storage — single file, persistent across restarts. Logged-in
# cookies live here. Gitignored.
STORAGE_STATE_PATH = Path(__file__).parent / "storage_state.json"

# Required CIPC prepaid balance for name reservation.
BALANCE_MIN = float(os.getenv("BIZPORTAL_BALANCE_MIN", "50"))

# Browser behaviour.
HEADLESS = os.getenv("BIZPORTAL_HEADLESS", "true").lower() == "true"
NAV_TIMEOUT_MS = int(os.getenv("BIZPORTAL_NAV_TIMEOUT_MS", "30000"))

# ─────────────────────────────────────────────────────────────────
# Logging
# ─────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO"),
    format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
)
log = logging.getLogger("bizportal-2")

# ─────────────────────────────────────────────────────────────────
# Pydantic models
# ─────────────────────────────────────────────────────────────────

class SearchRequest(BaseModel):
    type: str = Field(pattern="^(id|ck)$")
    query: str = Field(min_length=1, max_length=64)

class ReserveRequest(BaseModel):
    names: List[str] = Field(min_length=1, max_length=4)
    dry_run: bool = False
    balance_min: Optional[float] = None

# ─────────────────────────────────────────────────────────────────
# App + browser lifecycle
# ─────────────────────────────────────────────────────────────────

app = FastAPI(title="BizPortal 2.0", version="2.0.0")
_playwright = None
_browser: Optional[Browser] = None
_context: Optional[BrowserContext] = None
_browser_lock = asyncio.Lock()

@app.on_event("startup")
async def on_startup():
    global _playwright, _browser, _context
    _playwright = await async_playwright().start()
    # Use the full Chromium (not the headless-shell variant) — the headless
    # shell can hit EPERM on macOS when Playwright tries to clean up after
    # the first request. The full Chromium is more stable across platforms.
    launch_kwargs = {
        "headless": HEADLESS,
        "args": [
            "--no-sandbox",
            "--disable-dev-shm-usage",
            "--disable-gpu",
        ],
    }
    if HEADLESS:
        # Force the full Chromium binary, not the headless-shell.
        launch_kwargs["channel"] = "chromium"
    _browser = await _playwright.chromium.launch(**launch_kwargs)
    _context = await _open_context()
    log.info("Service ready. headless=%s storage_state=%s", HEADLESS, STORAGE_STATE_PATH.exists())

@app.on_event("shutdown")
async def on_shutdown():
    if _context:
        await _context.close()
    if _browser:
        await _browser.close()
    if _playwright:
        await _playwright.stop()

async def _open_context() -> BrowserContext:
    if STORAGE_STATE_PATH.exists():
        try:
            ctx = await _browser.new_context(storage_state=str(STORAGE_STATE_PATH))
            log.info("Loaded storage state from %s", STORAGE_STATE_PATH)
            return ctx
        except Exception as e:
            log.warning("Could not load storage state: %s — starting fresh", e)
    return await _browser.new_context()

def _current_context() -> BrowserContext:
    if not _context:
        raise RuntimeError("Browser context not initialised")
    return _context

# ─────────────────────────────────────────────────────────────────
# Health
# ─────────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {
        "ok": True,
        "service": "bizportal-2",
        "playwright_ready": _browser is not None and _context is not None,
        "logged_in_cached": STORAGE_STATE_PATH.exists(),
        "headless": HEADLESS,
    }

# ─────────────────────────────────────────────────────────────────
# Helpers — login, wait, parsing
# ─────────────────────────────────────────────────────────────────

async def _save_storage_state() -> None:
    state = await _context.storage_state()
    STORAGE_STATE_PATH.write_text(json.dumps(state, indent=2))
    log.info("Saved storage state to %s", STORAGE_STATE_PATH)

async def _is_logged_in(page: Page) -> bool:
    try:
        return await page.locator(LOGIN_ID).count() == 0
    except Exception:
        return False

async def _ensure_logged_in(page: Page, log_list: List[str]) -> None:
    """Reuse stored session if valid; otherwise log in once."""
    if await _is_logged_in(page):
        log_list.append("Already logged in (cached session)")
        return

    log.info("Logging in…")
    log_list.append("Navigating to login")
    await page.goto(LOGIN_URL, wait_until="domcontentloaded", timeout=NAV_TIMEOUT_MS)

    if await _is_logged_in(page):
        log_list.append("Already logged in (after navigation)")
        return

    id_no = os.getenv("BIZPORTAL_ID_NO")
    password = os.getenv("BIZPORTAL_PASSWORD")
    if not id_no or not password:
        raise RuntimeError("BIZPORTAL_ID_NO / BIZPORTAL_PASSWORD not set in environment")

    await page.locator(f"#{LOGIN_ID}").fill(id_no)
    await page.locator(f"#{LOGIN_PASSWORD}").fill(password)
    await page.locator(f"#{LOGIN_BUTTON}").click()

    # Wait for the login form to disappear.
    try:
        await page.locator(f"#{LOGIN_ID}").wait_for(state="detached", timeout=30000)
    except PWTimeout:
        raise RuntimeError("Login failed — login form still present after 30s")

    await _save_storage_state()
    log_list.append("Logged in and cached session")
    log.info("Logged in")

async def _parse_balance(page: Page) -> Optional[float]:
    await page.goto(PROFILE_URL, wait_until="domcontentloaded", timeout=NAV_TIMEOUT_MS)
    await page.wait_for_timeout(1500)
    text = await page.locator("body").inner_text()
    matches = re.findall(r"CIPC Prepaid Account Balance[:\s]*R\s*([0-9][0-9.,]*)", text, re.IGNORECASE)
    if not matches:
        return None
    raw = matches[0].replace(".", "").replace(",", ".")
    try:
        return float(raw)
    except ValueError:
        return None

async def _select_first(page: Page, candidates: List[str], kind: str) -> str:
    """Return the first selector that resolves. kind is for error msgs."""
    for sel in candidates:
        try:
            count = await page.locator(sel).count()
            if count > 0:
                return sel
        except Exception:
            continue
    raise RuntimeError(
        f"Could not find {kind} on {page.url}. "
        f"Tried: {candidates}. "
        f"Update the *_CANDIDATES constants in bizportal_service.py."
    )


async def _parse_results_table(page: Page) -> List[Dict[str, Any]]:
    """Probe-based parse that knows the actual BizPortal layouts.

    Two distinct result shapes:

    CK search (dropdown=EntNo): ONE company shown on the page itself.
      Labels: cntMain_lblEntNo/EntName/EntType/EntStatus hold the
      company info. Directors live in a sub-table with columns
      `ID Number | Name(s) | Surname | Type | Status`.

    ID search (dropdown=IDNo): a GridView of companies.
      `cntMain_gdvEnterprises` has columns Enterprise Number, Name,
      Status, Director Status, [hidden Status], View. Each row has
      a View button that triggers a postback to load that company's
      directors. We click through each one (up to a cap) and parse
      the directors from the re-rendered page.
    """
    # Detect which layout we're in.
    has_company_labels = await page.locator(f"#{BIZPROFILE_LBL_ENT_NO}").count() > 0
    has_grid = await page.locator(f"#{ID_SEARCH_GRID}").count() > 0

    if has_company_labels:
        return [await _parse_company_layout(page)]
    if has_grid:
        return await _parse_grid_layout(page)

    log_list: List[str] = []
    raise RuntimeError(
        "Unrecognised result layout. Neither company labels nor "
        f"the GridView were found on {page.url}."
    )


async def _parse_company_layout(page: Page) -> Dict[str, Any]:
    """CK search result — single company, directors in a sub-table."""
    async def lbl(suffix: str) -> str:
        loc = page.locator(f"#cntMain_lbl{suffix}")
        if await loc.count() == 0:
            return ""
        return (await loc.first.inner_text()).strip()

    company = {
        "enterprise_no": await lbl("EntNo"),
        "enterprise_name": await lbl("EntName"),
        "enterprise_type": await lbl("EntType"),
        "enterprise_status": await lbl("EntStatus"),
        "directors": [],
    }

    # Directors table — find the table whose header row contains
    # "ID Number". Verified ordering: ID Number | Name(s) | Surname
    # | Type | Status.
    tables = await page.locator("table").all()
    for t in tables:
        try:
            headers = await t.locator("th").all()
            header_text = [(await h.inner_text()).strip() for h in headers]
        except Exception:
            continue
        if not header_text:
            continue
        if any("ID Number" in h for h in header_text):
            cols = {"id": None, "names": None, "surname": None, "type": None, "status": None}
            for i, h in enumerate(header_text):
                hl = h.lower()
                if "id number" in hl:
                    cols["id"] = i
                elif "name(s)" in hl or hl.startswith("name"):
                    cols["names"] = i
                elif "surname" in hl:
                    cols["surname"] = i
                elif "type" in hl:
                    cols["type"] = i
                elif "status" in hl:
                    cols["status"] = i
            for row in await t.locator("tbody tr").all():
                cells = await row.locator("td").all()
                if not cells:
                    continue
                cell_text = []
                for c in cells:
                    try:
                        cell_text.append((await c.inner_text()).strip())
                    except Exception:
                        cell_text.append("")
                if all(not x for x in cell_text):
                    continue
                def at(k: str) -> str:
                    i = cols.get(k)
                    if i is None or i >= len(cell_text):
                        return ""
                    return cell_text[i]
                company["directors"].append({
                    "id_number": at("id"),
                    "names": at("names"),
                    "surname": at("surname"),
                    "status": at("status"),
                    "type": at("type") or "Director",
                })
            break

    return company


async def _parse_grid_layout(page: Page) -> List[Dict[str, Any]]:
    """ID search result — list of companies in a GridView.

    Returns the company list immediately. Directors are NOT fetched
    here — the caller can request them per-company via the
    /directors/<enterprise_no> endpoint. This keeps the search
    fast even when the ID has many director roles.
    """
    companies: List[Dict[str, Any]] = []
    rows = await page.locator(f"#{ID_SEARCH_GRID} tbody tr").all()
    for row in rows:
        cells = await row.locator("td").all()
        if len(cells) < 4:
            continue
        cell_text = []
        for c in cells:
            try:
                cell_text.append((await c.inner_text()).strip())
            except Exception:
                cell_text.append("")
        if not cell_text[0]:
            continue
        companies.append({
            "enterprise_no": cell_text[0],
            "enterprise_name": cell_text[1] if len(cell_text) > 1 else "",
            "enterprise_status": cell_text[2] if len(cell_text) > 2 else "",
            "director_status": cell_text[3] if len(cell_text) > 3 else "",
            "directors": [],
        })
    return companies


async def _fetch_directors_for_company(page: Page, enterprise_no: str) -> List[Dict[str, Any]]:
    """ID search follow-up — do a separate CK search for the given
    enterprise number and return its directors.

    Clicking the View button on the GridView navigates to the
    company details page (the GridView is gone afterwards), so we
    can't click multiple in a row. Doing a fresh CK search is the
    simplest reliable path.
    """
    await page.goto(BIZPROFILE_URL, wait_until="domcontentloaded", timeout=NAV_TIMEOUT_MS)
    await page.wait_for_timeout(1500)
    await page.locator(f"#{BIZPROFILE_DROPDOWN}").select_option("EntNo")
    try:
        await page.wait_for_load_state("domcontentloaded", timeout=15000)
    except Exception:
        pass
    await page.wait_for_timeout(1500)
    await page.locator(f"#{BIZPROFILE_INPUT}").fill(enterprise_no)
    await page.locator(f"#{BIZPROFILE_BUTTON}").click()
    await page.wait_for_load_state("domcontentloaded", timeout=NAV_TIMEOUT_MS)
    await page.wait_for_timeout(2500)
    company = await _parse_company_layout(page)
    return company["directors"]

# ─────────────────────────────────────────────────────────────────
# POST /search
# ─────────────────────────────────────────────────────────────────

@app.post("/search")
async def search(req: SearchRequest):
    log_list: List[str] = []
    async with _browser_lock:
        page = await _context.new_page() if _context else None
        if not page:
            raise HTTPException(503, "Browser not ready")
        try:
            await _ensure_logged_in(page, log_list)

            option = BIZPROFILE_DROPDOWN_OPTIONS.get(req.type)
            if not option:
                return JSONResponse(
                    {"ok": False, "error": f"Unknown search type: {req.type}", "log": log_list},
                    status_code=400,
                )

            log_list.append(f"GET {BIZPROFILE_URL}")
            await page.goto(BIZPROFILE_URL, wait_until="domcontentloaded", timeout=NAV_TIMEOUT_MS)
            await page.wait_for_timeout(1500)

            log_list.append(f"Select dropdown: {option}")
            await page.locator(f"#{BIZPROFILE_DROPDOWN}").select_option(option)
            # Dropdown change triggers a postback — wait for the page to settle.
            try:
                await page.wait_for_load_state("domcontentloaded", timeout=15000)
            except Exception:
                pass
            await page.wait_for_timeout(1500)

            log_list.append(f"Fill query: {req.query}")
            await page.locator(f"#{BIZPROFILE_INPUT}").fill(req.query)
            await page.locator(f"#{BIZPROFILE_BUTTON}").click()
            await page.wait_for_load_state("domcontentloaded", timeout=NAV_TIMEOUT_MS)
            await page.wait_for_timeout(3000)

            companies = await _parse_results_table(page)
            log_list.append(f"Parsed {len(companies)} companies")

            # For ID search, optionally click-through to fetch directors.
            if req.type == "id" and companies:
                capped = companies[:MAX_ID_SEARCH_DIRECTOR_FETCH]
                skipped = len(companies) - len(capped)
                log_list.append(f"Fetching directors for {len(capped)} companies (cap={MAX_ID_SEARCH_DIRECTOR_FETCH})")
                for c in capped:
                    directors = await _fetch_directors_for_company(page, c["enterprise_no"])
                    c["directors"] = directors
                    log_list.append(f"  {c['enterprise_no']}: {len(directors)} directors")
                if skipped:
                    log_list.append(f"Skipped {skipped} additional companies to keep search fast")

            return {"ok": True, "query": req.query, "type": req.type, "results": companies, "log": log_list}

        except Exception as e:
            log.exception("Search failed")
            log_list.append(f"ERROR: {e}")
            return JSONResponse({"ok": False, "error": str(e), "log": log_list}, status_code=500)
        finally:
            await page.close()

# ─────────────────────────────────────────────────────────────────
# POST /reserve
# ─────────────────────────────────────────────────────────────────

@app.post("/reserve")
async def reserve(req: ReserveRequest):
    log_list: List[str] = []
    balance_before: Optional[float] = None
    threshold = req.balance_min if req.balance_min is not None else BALANCE_MIN
    balance_before = None
    step = "init"

    async with _browser_lock:
        page = await _context.new_page() if _context else None
        if not page:
            raise HTTPException(503, "Browser not ready")
        try:
            names = [n for n in (req.names or []) if n and n.strip()]
            if not names:
                raise RuntimeError("No proposed names to apply")
            log_list.append(f"Names: {names}")

            step = "login"
            await _ensure_logged_in(page, log_list)

            step = "balance_before"
            balance_before = await _parse_balance(page)
            log_list.append(f"Balance BEFORE: {balance_before}")
            if balance_before is None:
                raise RuntimeError("Could not read CIPC Prepaid Account Balance")
            if balance_before <= threshold:
                raise RuntimeError(
                    f"Insufficient balance: R {balance_before:.2f} is not > R {threshold:.2f}. "
                    f"Top up the CIPC Prepaid Account before submitting."
                )

            step = "open_name_reservation"
            log_list.append(f"GET {NAME_RESERVE_URL}")
            await page.goto(NAME_RESERVE_URL, wait_until="domcontentloaded", timeout=NAV_TIMEOUT_MS)
            await page.wait_for_timeout(1500)

            step = "fill_names"
            for i, name in enumerate(names[:4]):
                sel = f"#{NAME_INPUTS[i]}"
                count = await page.locator(sel).count()
                if count == 0:
                    raise RuntimeError(f"Name field {sel} not found on name_reserve page")
                await page.locator(sel).fill(name)
            log_list.append(f"Filled {len(names[:4])} name boxes")

            step = "continue_1"
            await page.locator(f"#{NAME_CONTINUE_1}").wait_for(state="visible", timeout=10000)
            await page.locator(f"#{NAME_CONTINUE_1}").click()
            log_list.append("Clicked first Continue")
            await page.wait_for_timeout(2500)

            if req.dry_run:
                log_list.append("DRY RUN — stopping before second Continue")
                return {
                    "ok": True,
                    "dry_run": True,
                    "balance_before": balance_before,
                    "status": "dry_run",
                    "log": log_list,
                    "message": "Dry run complete. Cost review screen reached. Re-run with dry_run=false to lodge.",
                }

            step = "continue_2"
            await page.locator(f"#{NAME_CONTINUE_2}").wait_for(state="visible", timeout=10000)
            await page.locator(f"#{NAME_CONTINUE_2}").click()
            log_list.append("Clicked second Continue")
            await page.wait_for_timeout(4000)

            return {
                "ok": True,
                "balance_before": balance_before,
                "status": "bizportal_done",
                "log": log_list,
                "message": "Submitted via Bizportal. The R50 fee will be debited from your CIPC Prepaid Account at a later stage. You will receive an SMS confirmation.",
            }

        except Exception as e:
            log_list.append(f"ERROR at step={step}: {e}")
            return JSONResponse(
                {"ok": False, "error": str(e), "step": step, "balance_before": balance_before, "log": log_list},
                status_code=500,
            )
        finally:
            await page.close()

# ─────────────────────────────────────────────────────────────────
# Entrypoint
# ─────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("bizportal_service:app", host="127.0.0.1", port=7001, log_level="info")