from __future__ import annotations

from typing import cast

from selenium.webdriver.common.by import By
from selenium.webdriver.remote.webdriver import WebDriver
from selenium.webdriver.remote.webelement import WebElement
from typing_extensions import TypedDict

from .schema import Section


class TableJson(TypedDict):
    title: str
    columns: list[str]
    values: list[dict[str, str]]


class SectionJson(TypedDict):
    section: str | None
    tables: list[TableJson]


def _headers_from_table(table: WebElement) -> list[str]:
    ths = table.find_elements(By.CSS_SELECTOR, "thead th")  # pyright: ignore[reportUnknownMemberType]
    result: list[str] = []
    for th in ths:
        links = th.find_elements(By.TAG_NAME, "a")  # pyright: ignore[reportUnknownMemberType]
        result.append(links[0].text if links else th.text)
    return result


def _rows_from_table(table: WebElement) -> list[list[str]]:
    rows: list[list[str]] = []
    for tr in table.find_elements(By.CSS_SELECTOR, "tbody tr"):  # pyright: ignore[reportUnknownMemberType]
        tds = tr.find_elements(By.TAG_NAME, "td")  # pyright: ignore[reportUnknownMemberType]
        rows.append([td.text for td in tds])
    return rows


def _table_json(table: WebElement) -> TableJson | None:
    columns = _headers_from_table(table)
    if not columns:
        return None
    rows = _rows_from_table(table)
    title = columns[0]

    values = []

    for row in rows:
        v = {}
        for c, r in zip(columns, row, strict=False):
            v[c] = r
        values.append(v)  # pyright: ignore[reportUnknownMemberType]

    return TableJson(
        {
            "title": title,
            "columns": columns,
            "values": values,
        }
    )


def sections_to_json(driver: WebDriver) -> list[Section]:
    seq = driver.find_elements(By.CSS_SELECTOR, "h3, table.eamwebgrid-table")

    sections: list[SectionJson] = []
    current: SectionJson = {"section": None, "tables": []}
    for el in seq:
        tag = el.tag_name.lower()
        if tag == "h3":
            if current["section"] is not None or current["tables"]:
                sections.append(current)
            current = {"section": el.text, "tables": []}
        else:
            table = _table_json(el)
            if table:
                current["tables"].append(table)

    if current["section"] is not None or current["tables"]:
        sections.append(current)

    # Schema is not yet complete, ignore validation
    return cast(list[Section], sections)
    # ta = TypeAdapter(list[Section])

    # return ta.validate_python(sections, by_alias=True)
