from __future__ import annotations

import os
from dataclasses import dataclass, field

from . import ast as AST
from .errors import SupRuntimeError


@dataclass
class IOHooks:
    stdin: str | None = None
    outputs: list[str] = field(default_factory=list)

    def read_input(self) -> str:
        if self.stdin is None:
            # Interactive input
            return input()
        # Pop first line from stdin buffer
        if "\n" in self.stdin:
            line, rest = self.stdin.split("\n", 1)
            self.stdin = rest
        else:
            line, self.stdin = self.stdin, None
        return line

    def write_output(self, text: str) -> None:
        self.outputs.append(text)


class Interpreter:
    def __init__(self) -> None:
        self.env: dict[str, object] = {}
        self.functions: dict[str, AST.FunctionDef] = {}
        self.module_cache: dict[str, dict[str, object]] = {}
        self.loading_modules: set[str] = set()
        self.last_result: object | None = None
        self.io = IOHooks()

    def run(self, program: AST.Program, *, stdin: str | None = None) -> str:
        self.io.stdin = stdin
        self.eval_program(program)
        return "".join(self.io.outputs)

    def eval_program(self, program: AST.Program) -> None:
        for stmt in program.statements:
            self.eval(stmt)

    def eval(self, node: AST.Node) -> object | None:
        if isinstance(node, AST.Assignment):
            value = self.eval(node.expr)
            self.env[node.name.lower()] = value
            self.last_result = value
            return value
        if isinstance(node, AST.Print):
            value = self.last_result if node.expr is None else self.eval(node.expr)
            self.io.write_output(f"{self._format_value(value)}\n")
            return value
        if isinstance(node, AST.Ask):
            val = self.io.read_input()
            self.env[node.name.lower()] = val
            self.last_result = val
            return val
        if isinstance(node, AST.If):
            cond_val = self._truthy(self.eval(node.cond) if node.cond is not None else self._compare(self.eval(node.left), node.op, self.eval(node.right)))  # type: ignore[arg-type]
            if cond_val:
                for s in node.body or []:
                    self.eval(s)
            else:
                for s in node.else_body or []:
                    self.eval(s)
            return None
        if isinstance(node, AST.While):
            while self._truthy(self.eval(node.cond)):
                for s in node.body:
                    self.eval(s)
            return None
        if isinstance(node, AST.ForEach):
            iterable = self.eval(node.iterable)
            try:
                iterator = list(iterable)  # type: ignore[arg-type]
            except Exception:
                raise SupRuntimeError(message="Target of for each is not iterable.")
            saved = self.env.get(node.var.lower())
            try:
                for item in iterator:
                    self.env[node.var.lower()] = item
                    for s in node.body:
                        self.eval(s)
            finally:
                if saved is None:
                    self.env.pop(node.var.lower(), None)
                else:
                    self.env[node.var.lower()] = saved
            return None
        if isinstance(node, AST.Repeat):
            count_val = self.eval(node.count_expr)
            try:
                iterations = int(count_val)  # type: ignore[arg-type]
            except Exception:
                raise SupRuntimeError(
                    message="Repeat count must be a number.",
                    line=getattr(node, "line", None),
                )
            for _ in range(iterations):
                for s in node.body:
                    self.eval(s)
            return None
        if isinstance(node, AST.ExprStmt):
            value = self.eval(node.expr)
            self.last_result = value
            return value
        if isinstance(node, AST.TryCatch):
            error: Exception | None = None
            try:
                for s in node.body:
                    self.eval(s)
            except Exception as e:  # catch Sup and general errors
                error = e
                if node.catch_body is not None:
                    if node.catch_name:
                        if isinstance(e, _SupThrown):
                            self.env[node.catch_name.lower()] = e.value
                        else:
                            self.env[node.catch_name.lower()] = str(e)
                    for s in node.catch_body:
                        self.eval(s)
                else:
                    # no catch: rethrow after finally
                    pass
            finally:
                if node.finally_body is not None:
                    for s in node.finally_body:
                        self.eval(s)
                if error is not None and node.catch_body is None:
                    raise error
            return None
        if isinstance(node, AST.Throw):
            val = self.eval(node.value)
            raise _SupThrown(val)
        if isinstance(node, AST.Import):
            ns = self._import_module(node.module)
            self.env[(node.alias or node.module).lower()] = ns
            return None
        if isinstance(node, AST.FromImport):
            ns = self._import_module(node.module)
            for name, alias in node.names:
                if name not in ns:
                    raise SupRuntimeError(
                        message=f"Module '{node.module}' has no symbol '{name}'."
                    )
                self.env[(alias or name).lower()] = ns[name]
            return None
        if isinstance(node, AST.FunctionDef):
            self.functions[node.name.lower()] = node
            return None
        if isinstance(node, AST.Return):
            # Signal a return using exception for simple control flow
            raise _ReturnSignal(self.eval(node.expr) if node.expr is not None else None)
        if isinstance(node, AST.Call):
            return self._call_function(node)
        # Collections and stdlib
        if isinstance(node, AST.MakeList):
            lst = [self.eval(it) for it in node.items]
            self.env["list"] = lst
            self.last_result = lst
            return lst
        if isinstance(node, AST.MakeMap):
            d: dict[object, object] = {}
            self.env["map"] = d
            self.last_result = d
            return d
        if isinstance(node, AST.Push):
            target = self.eval(node.target)
            if not isinstance(target, list):
                raise SupRuntimeError(message="Push target must be a list.")
            target.append(self.eval(node.item))
            self.last_result = target
            return target
        if isinstance(node, AST.Pop):
            target = self.eval(node.target)
            if not isinstance(target, list):
                raise SupRuntimeError(message="Pop target must be a list.")
            val = target.pop()
            self.last_result = val
            return val
        if isinstance(node, AST.GetKey):
            target = self.eval(node.target)
            key = self.eval(node.key)
            if isinstance(target, list):
                try:
                    idx = int(self._num(key))
                except Exception:
                    raise SupRuntimeError(message="List index must be a number.")
                try:
                    val = target[idx]
                except Exception:
                    raise SupRuntimeError(message="List index out of range.")
                self.last_result = val
                return val
            if isinstance(target, dict):
                val = target.get(key)
                self.last_result = val
                return val
            raise SupRuntimeError(message="Get target must be a list or map.")
        if isinstance(node, AST.SetKey):
            target = self.eval(node.target)
            key = self.eval(node.key)
            val = self.eval(node.value)
            if not isinstance(target, dict):
                raise SupRuntimeError(message="Set target must be a map.")
            target[key] = val
            self.last_result = target
            return target
        if isinstance(node, AST.DeleteKey):
            target = self.eval(node.target)
            key = self.eval(node.key)
            if not isinstance(target, dict):
                raise SupRuntimeError(message="Delete target must be a map.")
            target.pop(key, None)
            self.last_result = target
            return target
        if isinstance(node, AST.Length):
            target = self.eval(node.target)
            length_value = len(target)  # type: ignore[arg-type]
            self.last_result = length_value
            return length_value
        if isinstance(node, AST.BuiltinCall):
            return self._eval_builtin(node)
        if isinstance(node, AST.Binary):
            left = self.eval(node.left)
            right = self.eval(node.right)
            try:
                if node.op in {"+", "-", "*"}:
                    lnum, lint = self._to_number(left)
                    rnum, rint = self._to_number(right)
                    if node.op == "+":
                        value = lnum + rnum
                    elif node.op == "-":
                        value = lnum - rnum
                    else:
                        value = lnum * rnum
                    if lint and rint and float(value).is_integer():
                        res = int(value)
                    else:
                        res = float(value)
                elif node.op == "/":
                    res = float(self._num(left)) / float(self._num(right))
                else:
                    raise SupRuntimeError(message=f"Unknown operator {node.op}.")
            except ZeroDivisionError:
                raise SupRuntimeError(
                    message="Division by zero.", line=getattr(node, "line", None)
                )
            self.last_result = res
            return res
        if isinstance(node, AST.Identifier):
            name = node.name.lower()
            # dotted access: module.symbol
            if "." in name:
                mod, sym = name.split(".", 1)
                if mod in self.env and isinstance(self.env[mod], dict):
                    ns = self.env[mod]
                    return ns.get(sym)
            if name in self.env:
                return self.env[name]
            # Allow implicit references to 'list' and 'map' if they were just created as last_result
            if name in {"list", "map"} and isinstance(self.last_result, (list, dict)):
                return self.last_result
            raise SupRuntimeError(
                message=f"Undefined variable '{node.name}'.",
                line=getattr(node, "line", None),
            )
        if isinstance(node, AST.String):
            return node.value
        if isinstance(node, AST.Number):
            return node.value
        if isinstance(node, AST.BoolBinary):
            if node.op == "and":
                left = self._truthy(self.eval(node.left))
                return left and self._truthy(self.eval(node.right))
            if node.op == "or":
                left = self._truthy(self.eval(node.left))
                return left or self._truthy(self.eval(node.right))
            raise SupRuntimeError(message=f"Unknown boolean operator {node.op}.")
        if isinstance(node, AST.NotOp):
            return not self._truthy(self.eval(node.expr))
        if isinstance(node, AST.Compare):
            return self._compare(self.eval(node.left), node.op, self.eval(node.right))
        raise SupRuntimeError(message=f"Unsupported AST node {type(node).__name__}.")

    def _compare(self, left: object, op: str | None, right: object) -> bool:  # type: ignore[override]
        if op == ">":
            return self._num(left) > self._num(right)
        if op == "<":
            return self._num(left) < self._num(right)
        if op == "==":
            return left == right
        if op == "!=":
            return left != right
        if op == ">=":
            return self._num(left) >= self._num(right)
        if op == "<=":
            return self._num(left) <= self._num(right)
        raise SupRuntimeError(message=f"Unknown relational operator {op}.")

    def _truthy(self, v: object) -> bool:
        return bool(v)

    def _num(self, v: object) -> float:
        if isinstance(v, (int, float)):
            return float(v)
        # attempt to parse string numbers for friendliness
        if isinstance(v, str):
            try:
                return float(v)
            except ValueError:
                pass
        raise SupRuntimeError(message=f"Expected a number, got {type(v).__name__}.")

    def _format_value(self, v: object) -> str:
        if isinstance(v, float):
            return str(v)
        return str(v)

    def _to_number(self, v: object) -> tuple[float, bool]:
        is_int = isinstance(v, int)
        num = self._num(v)
        return num, is_int

    def _eval_builtin(self, node: AST.BuiltinCall) -> object:
        name = node.name
        if name == "now":
            import datetime as _dt

            res = _dt.datetime.now().isoformat()
            self.last_result = res
            return res
        if name == "read_file":
            path = str(self.eval(node.args[0]))
            with open(path, encoding="utf-8") as f:
                res = f.read()
            self.last_result = res
            return res
        if name == "write_file":
            path = str(self.eval(node.args[0]))
            data = str(self.eval(node.args[1]))
            with open(path, "w", encoding="utf-8") as f:
                f.write(data)
            self.last_result = True
            return True
        if name == "json_parse":
            import json as _json

            s = str(self.eval(node.args[0]))
            res = _json.loads(s)
            self.last_result = res
            return res
        if name == "json_stringify":
            import json as _json

            v = self.eval(node.args[0])
            res = _json.dumps(v)
            self.last_result = res
            return res
        if name == "min":
            a = self.eval(node.args[0])
            b = self.eval(node.args[1])
            res = min(self._num(a), self._num(b))
            self.last_result = float(res)
            return float(res)
        if name == "max":
            a = self.eval(node.args[0])
            b = self.eval(node.args[1])
            res = max(self._num(a), self._num(b))
            self.last_result = float(res)
            return float(res)
        if name == "floor":
            import math

            a = self._num(self.eval(node.args[0]))
            res = math.floor(a)
            self.last_result = float(res)
            return float(res)
        if name == "ceil":
            import math

            a = self._num(self.eval(node.args[0]))
            res = math.ceil(a)
            self.last_result = float(res)
            return float(res)
        if name == "trim":
            s = str(self.eval(node.args[0]))
            res = s.strip()
            self.last_result = res
            return res
        if name == "contains":
            s = self.eval(node.args[0])
            sub = self.eval(node.args[1])
            if isinstance(s, list):
                res = any(item == sub for item in s)
            else:
                res = str(sub) in str(s)
            self.last_result = res
            return res
        if name == "join":
            sep = str(self.eval(node.args[0]))
            lst = self.eval(node.args[1])
            if not isinstance(lst, list):
                raise SupRuntimeError(message="join expects a list.")
            res = sep.join(str(x) for x in lst)
            self.last_result = res
            return res
        if name == "power":
            a = self._num(self.eval(node.args[0]))
            b = self._num(self.eval(node.args[1]))
            res = float(a) ** float(b)
            self.last_result = res
            return res
        if name == "sqrt":
            import math

            a = self._num(self.eval(node.args[0]))
            res = math.sqrt(float(a))
            self.last_result = res
            return res
        if name == "abs":
            a = self.eval(node.args[0])
            if isinstance(a, (int, float)):
                res = abs(a)
            else:
                res = abs(self._num(a))
            self.last_result = float(res)
            return float(res)
        if name == "upper":
            s = str(self.eval(node.args[0]))
            res = s.upper()
            self.last_result = res
            return res
        if name == "lower":
            s = str(self.eval(node.args[0]))
            res = s.lower()
            self.last_result = res
            return res
        if name == "concat":
            a = str(self.eval(node.args[0]))
            b = str(self.eval(node.args[1]))
            res = a + b
            self.last_result = res
            return res
        raise SupRuntimeError(message=f"Unknown builtin {name}.")

    def _call_function(self, node: AST.Call) -> object:
        # module-qualified call mm.square
        name = node.name.lower()
        if "." in name:
            mod, sym = name.split(".", 1)
            if mod in self.env and isinstance(self.env[mod], dict):
                target = self.env[mod].get(sym)
                if isinstance(target, AST.FunctionDef):
                    return self._call_fn_def(target, node.args)
                raise SupRuntimeError(
                    message=f"Undefined function '{node.name}'.",
                    line=getattr(node, "line", None),
                )
        # direct function from env via from-import
        if name in self.env and isinstance(self.env[name], AST.FunctionDef):
            return self._call_fn_def(self.env[name], node.args)  # type: ignore[arg-type]
        if name not in self.functions:
            raise SupRuntimeError(
                message=f"Undefined function '{node.name}'.",
                line=getattr(node, "line", None),
            )
        fn = self.functions[name]
        if len(node.args) != len(fn.params):
            raise SupRuntimeError(
                message=f"Function '{fn.name}' expects {len(fn.params)} argument(s) but got {len(node.args)}."
            )
        # Evaluate args
        arg_vals = [self.eval(a) for a in node.args]
        # New scope
        saved_env = self.env.copy()
        try:
            self.env = self.env.copy()
            for pname, pval in zip(fn.params, arg_vals):
                self.env[pname.lower()] = pval
            # Execute body
            ret_val: object | None = None
            try:
                for s in fn.body:
                    self.eval(s)
            except _ReturnSignal as r:
                ret_val = r.value
            self.last_result = ret_val
            # Ensure division/float semantics are visible when printed via call in print
            if isinstance(ret_val, (int, float)) and not isinstance(ret_val, bool):
                # Keep numeric as float if any arithmetic implied result is float
                # For simplicity, cast ints to float only when returned? Tests expect 12.0 here
                return float(ret_val)
            return ret_val
        finally:
            self.env = saved_env

    def _call_fn_def(self, fn: AST.FunctionDef, arg_nodes: list[AST.Node]) -> object:
        if len(arg_nodes) != len(fn.params):
            raise SupRuntimeError(
                message=f"Function '{fn.name}' expects {len(fn.params)} argument(s) but got {len(arg_nodes)}."
            )
        arg_vals = [self.eval(a) for a in arg_nodes]
        saved_env = self.env.copy()
        try:
            self.env = self.env.copy()
            for pname, pval in zip(fn.params, arg_vals):
                self.env[pname.lower()] = pval
            ret_val: object | None = None
            try:
                for s in fn.body:
                    self.eval(s)
            except _ReturnSignal as r:
                ret_val = r.value
            self.last_result = ret_val
            if isinstance(ret_val, (int, float)) and not isinstance(ret_val, bool):
                return float(ret_val)
            return ret_val
        finally:
            self.env = saved_env

    def _import_module(self, module: str) -> dict[str, object]:
        key = module.lower()
        if key in self.module_cache:
            return self.module_cache[key]
        if key in self.loading_modules:
            raise SupRuntimeError(
                message=f"Circular import detected for module '{module}'."
            )
        # Resolve path
        search_paths = [os.getcwd()]
        env_path = os.environ.get("SUP_PATH")
        if env_path:
            search_paths = env_path.split(os.pathsep) + search_paths
        path = None
        for base in search_paths:
            candidate = os.path.join(base, f"{module}.sup")
            if os.path.exists(candidate):
                path = candidate
                break
        if path is None:
            raise SupRuntimeError(message=f"Cannot find module '{module}'.")
        # Load file and execute in fresh interpreter sharing module cache
        with open(path, encoding="utf-8") as f:
            src = f.read()
        from .parser import Parser

        parser = Parser()
        program = parser.parse(src)
        self.loading_modules.add(key)
        try:
            child = Interpreter()
            child.module_cache = self.module_cache  # share cache
            child.loading_modules = self.loading_modules
            child.run(program)
        finally:
            self.loading_modules.discard(key)
        # Export top-level env and functions
        ns: dict[str, object] = {}
        ns.update(child.env)
        for name, fn in child.functions.items():
            ns[name] = fn
        self.module_cache[key] = ns
        return ns


class _ReturnSignal(Exception):
    def __init__(self, value: object | None) -> None:
        self.value = value


class _SupThrown(Exception):
    def __init__(self, value: object) -> None:
        self.value = value
        super().__init__(str(value))
