from pydantic import BaseModel
from pydantic.config import ConfigDict
import pandas as pd
from textual.coordinate import Coordinate
from textual.app import App, ComposeResult
from textual.widgets import Header, Footer, DataTable, Static, TabbedContent, TabPane, Label
from textual.containers import Vertical, VerticalScroll
from rich.text import Text
from rich.style import Style
from typing import Self, Optional

from smtconfig.utils import free_vars

from .config import ConfigInstance, FieldInst, FieldStatus, Rule, SMTResult
from icecream import ic

STATUS_COLORS = {
    FieldStatus.PROVIDED: "cyan",
    FieldStatus.SOLVED: "green",
    FieldStatus.CONFLICT: "red",
    FieldStatus.UNCERTAIN: "yellow",
    FieldStatus.UNSOLVED: "magenta",
}

Label.link_style = Style(underline=True)


class ArbTypesModel(BaseModel):
    model_config: ConfigDict = ConfigDict(arbitrary_types_allowed=True)


class ViewerDocument(ArbTypesModel):
    name: str
    config_inst: ConfigInstance
    df: pd.DataFrame

    @classmethod
    def prep_document(cls, name: str, config_inst: ConfigInstance) -> Self:
        df = config_inst.build_table_dataframe()
        return cls(name=name, config_inst=config_inst, df=df)


class ConfigViewer(App):
    CSS = """
    #status {
        padding: 1;
        height: 3;
    }
    #bottom-pane {
        dock: bottom;
        height: 50%;
    }
    """

    def __init__(self,
                 config_inst: ConfigInstance,
                 title: Optional[str] = None):
        super().__init__()
        self.title = title
        self.sub_title = "ConfigViewer"
        self.config_inst = config_inst
        self.status_header = Static("kek", markup=True, id="status")
        self.df = self.config_inst.build_table_dataframe()

    def compose(self) -> ComposeResult:
        yield Header()
        yield Footer()
        with Vertical():
            yield self.status_header
            with VerticalScroll():
                yield DataTable(id="fields",
                                cursor_type="row")  # self.field_table
            with TabbedContent(id="bottom-pane"):
                with TabPane("Rules", id="rules-pane"):
                    yield Label("Rules", markup=True, id="rules")
                with TabPane("Versions", id="versions-pane"):
                    yield Static("Versions", markup=True, id="versions")
                with TabPane("Logs", id="logs-pane"):
                    yield Static("Logs", markup=True, id="logs")

    def _status_markup(self):
        s = self.config_inst.status
        if s.sat():
            return f"Status: [green]{s}[/green]"
        else:
            return f"Status: [red]{s}[/red]"

    def _fmt_field_status(self, st: FieldInst):
        c = STATUS_COLORS.get(st.status, "white")
        return f"[{c}]{str(st)}[/{c}]"

    def _find_field_row(self, field_name: str):
        data = self.df[self.df["Field"] == field_name]
        if len(data) == 0:
            return None
        assert len(data) == 1
        return next(iter(data.itertuples(index=False)))

    def on_mount(self):
        field_table = self.query_one(DataTable)
        field_table.add_columns("Field", "Value", "Type", "Status")
        for row in self.df.itertuples():
            field_table.add_row(row.Field, str(row.Value), row.Type,
                                self._fmt_field_status(row.Status))
        self.status_header.content = self._status_markup()

    def _format_rule_expr(self, rule: Rule) -> Text:
        rule_expr = str(rule.expr)
        fields = [str(var) for var in free_vars(rule.expr)]
        out = Text()
        i = 0
        n = len(rule_expr)

        fields_sorted = sorted(fields, key=len, reverse=True)

        while i < n:
            matched = False
            for f in fields_sorted:
                if rule_expr.startswith(f, i):
                    # raise Exception(
                    #     STATUS_COLORS[self._find_field_row(f).Status.status])
                    if not self._find_field_row(f):
                        continue
                    seg = Text(
                        f,
                        STATUS_COLORS[self._find_field_row(f).Status.status])
                    seg = seg.on(click=f"app.highlight_field('{f}')")
                    out.append(seg)
                    i += len(f)
                    matched = True
                    break
            if not matched:
                out.append(rule_expr[i])
                i += 1
        return out

    def action_highlight_field(self, field_name: str):
        idx = self.df.index[self.df["Field"] == field_name].tolist()
        assert len(idx) == 1
        idx = idx[0]
        field_table = self.query_one(DataTable)
        field_table.cursor_coordinate = Coordinate(row=idx, column=0)
        # field_table.

    def _refresh_rules(self, new_field: int):
        rules_info: Label = self.query_one("#rules", Label)
        # self.app.highligh_field("EEEEEEEe")
        field_name = self.df.loc[new_field]["Field"]
        rules = self.config_inst.config.rules_on_fields([field_name])

        def _get_rule_expr(rule: Rule):
            parts = [Text(rule.name), self._format_rule_expr(rule)]
            res = Text(": ").join(parts)
            return res

        if len(rules) > 0:
            rule_exprs = [_get_rule_expr(rule) for rule in rules]
            rules_info.content = (
                Text(f"'{field_name}' is involved in rules:\n", ) +
                Text('\n').join(
                    Text(f'{i}) ') + expr
                    for i, expr in enumerate(rule_exprs, start=1)))
        else:
            rules_info.content = f"'{field_name}' is not involved in any rules"

    def on_key(self, event):
        field_table = self.query_one(DataTable)
        if event.key in ("q", ):
            self.exit()
        elif event.key == "j":
            field_table.action_cursor_down()
        elif event.key == "k":
            field_table.action_cursor_up()

    def on_data_table_row_highlighted(self,
                                      event: DataTable.RowHighlighted) -> None:
        match event.data_table.id:
            case "fields":
                self._refresh_rules(event.cursor_row)
            case _:
                raise Exception(
                    f"Unexpected datatable '{event.data_table.id}'")
