from __future__ import annotations

import sys
from pathlib import Path
import json

import importlib.util
import os

import pyarrow as pa
import pyarrow.parquet as pq
import pprint
import numpy as np
import numpy.typing as npt

import networkx as nx

from automata.fa.dfa import DFA


_build_dir = Path(__file__).resolve().parent / "build"

if str(_build_dir) not in sys.path:
    sys.path.insert(0, str(_build_dir))

try:
    from . import _am_fast
    from ._am_fast import *
except ImportError as e:
    print(f"\n--- DEBUG START ---")
    print(f"Error: {e}")
    print(f"Current Directory: {os.getcwd()}")
    print(f"File Location: {__file__}")
    # This helps see if the .so file is actually in the folder
    print(f"Contents of this folder: {os.listdir(os.path.dirname(__file__))}")
    print(f"--- DEBUG END ---\n")
    raise e

# from ._am_fast import generate_cpp, block_entropy_convergence_cpp, strongly_connected_components_cpp, minify_dfa_cpp
from .json_utils import save_json

def block_entropy_convergence(
    h_mu: float,
    n_states: int,
    n_symbols : int,
    convergence_tol: float,
    precision: float,
    eps: float,
    branches: list[tuple[float, list[float]]],
    trans: list[list[tuple[int, float, int]]],
    max_branches: int = 30_000_000
) -> any :
    return block_entropy_convergence_cpp(
        h_mu            = float( h_mu ),
        n_states        = n_states,
        n_symbols       = n_symbols,
        convergence_tol = convergence_tol,
        precision       = float(precision),
        eps             = eps,
        branches        = branches,
        trans           = trans,
        max_branches    = max_branches
    )

def strongly_connected_components( T ) :
    return strongly_connected_components_cpp( T )

def generate_data(
    n_gen        : int,
    start_state  : int,
    transitions  : list[list[tuple[int, float, int]]],
    alphabet     : list[str],
    include_states: bool = False,
    random_seed : int = 42 ) -> dict[str, npt.NDArray]:

    n_states  = len(transitions)
    n_symbols = len(alphabet)

    if n_states > 65536:
        raise ValueError(f"Max states supported is 65536, got {n_states}")
    if n_symbols > 65536:
        raise ValueError(f"Max alphabet size supported is 65536, got {n_symbols}")

    symbol_indices, state_indices = generate_cpp(
        n_gen             = n_gen,
        start_state_index = start_state,
        transitions       = transitions,
        include_states    = include_states,
        random_seed       = random_seed
    )

    res: dict[str, npt.NDArray] = {
        "symbol_index": np.from_dlpack(symbol_indices)
    }
    if include_states:
        res["state_index"] = np.from_dlpack(state_indices)

    return res

def generate_data(
    n_gen          : int,
    start_state    : int,
    transitions    : list[list[tuple[int, float, int]]],
    alphabet       : list[str],
    include_states : bool = False,
    random_seed    : int  = 42,
) -> dict[str, npt.NDArray]:

    n_states  = len(transitions)
    n_symbols = len(alphabet)

    if n_states > 65536:
        raise ValueError(f"Max states supported is 65536, got {n_states}")
    if n_symbols > 65536:
        raise ValueError(f"Max alphabet size supported is 65536, got {n_symbols}")

    symbol_indices, state_indices = generate_cpp(
        n_gen             = n_gen,
        start_state_index = start_state,
        transitions       = transitions,
        include_states    = include_states,
        random_seed       = random_seed,
    )

    res: dict[str, npt.NDArray] = {
        "symbol_index": np.from_dlpack(symbol_indices)
    }
    if include_states:
        res["state_index"] = np.from_dlpack(state_indices)

    return res


def save_data(
    data             : dict[str, any],
    file_prefix      : str,
    alphabet         : list[str],
    n_states         : int,
    start_state      : int,
    random_seed      : int,
    machine_metadata : dict[str, any] | None = None,
) -> None:

    def _uint_type(n: int) -> pa.DataType:
        return pa.uint8() if n <= 256 else pa.uint16()

    sym_type   = _uint_type(len(alphabet))
    state_type = _uint_type(n_states)

    columns: dict[str, pa.Array] = {
        "symbol_index": pa.array(data["symbol_index"]).cast(sym_type)
    }

    if "state_index" in data:
        state_data  = data["state_index"]
        final_state = int(state_data[-1])
        columns["state_index"] = pa.array(state_data[:-1]).cast(state_type)

    iso_shifts = data.get("isomorphic_shifts", {})
    iso_final_states: dict[int, int] = {}

    for shift, shifted in iso_shifts.items():
        columns[f"symbol_index_isoshift_{shift}"] = pa.array(shifted["symbol_index"]).cast(sym_type)
        if "state_index" in shifted:
            iso_state_data                             = shifted["state_index"]
            iso_final_states[f"{shift}"]                    = int(iso_state_data[-1])
            columns[f"state_index_isoshift_{shift}"]   = pa.array(iso_state_data[:-1]).cast(state_type)

    parquet_meta: dict[str, any] = {
        "alphabet"          : alphabet,
        "machine_metadata"  : machine_metadata or {},
        "start_state"       : start_state,
        "random_seed"       : random_seed,
        "isomorphic_shifts" : sorted(iso_shifts.keys()),
    }

    if "state_index" in data:
        parquet_meta["final_state"] = final_state
    if iso_final_states:
        parquet_meta["isomorphic_final_states"] = iso_final_states

    table = pa.table(columns)
    table = table.replace_schema_metadata({
        **(table.schema.metadata or {}),
        "am_metadata": json.dumps(parquet_meta),
    })

    pq.write_table(table, f"{file_prefix}.parquet")
    save_json(parquet_meta, f"{file_prefix}.json")


def minify_cpp(dfa: DFA, retain_names: bool = True) -> DFA:
    state_list = list(dfa.states)
    state_idx = {s: i for i, s in enumerate(state_list)}
    symbol_list = sorted(dfa.input_symbols)
    symbol_idx = {sym: i for i, sym in enumerate(symbol_list)}

    adj = [[] for _ in range(len(state_list))]
    for state, paths in dfa.transitions.items():
        u = state_idx[state]
        for sym, nxt in paths.items():
            v = state_idx.get(nxt, -1)
            if v != -1:
                adj[u].append((symbol_idx[sym], v))

    is_final = [s in dfa.final_states for s in state_list]
    init_idx = state_idx[dfa.initial_state]

    # result names in res now match the C++ struct fields exactly
    res = minify_dfa_cpp(len(state_list), adj, init_idx, is_final)

    if res.is_empty_language:
        return dfa.__class__.empty_language(dfa.input_symbols)

    if retain_names:
        class_members = {}
        for old_idx, nc in enumerate(res.eq_class):
            if nc >= 0:
                class_members.setdefault(nc, []).append(state_list[old_idx])
        class_map = {nc: frozenset(m) for nc, m in class_members.items()}
    else:
        class_map = {nc: nc for nc in range(res.n_classes)}

    new_states = set(class_map.values())
    new_initial = class_map[res.new_initial]
    new_final = {class_map[c] for c, f in enumerate(res.class_is_final) if f}

    new_trans = {}
    for nc in range(res.n_classes):
        new_trans[class_map[nc]] = {
            symbol_list[trans[0]]: class_map[trans[1]]
            for trans in res.class_trans[nc]
        }

    return dfa.__class__(
        states=new_states,
        input_symbols=dfa.input_symbols,
        transitions=new_trans,
        initial_state=new_initial,
        final_states=new_final,
        allow_partial=any(len(t) < len(symbol_list) for t in new_trans.values())
    )