from __future__ import annotations

import re
from typing import Any
from .featurescript_lexer import Token, lex
from .ir import Call, FeatureIR, ModelIR, SketchIR


def _string(value: Any) -> Any:
    if isinstance(value, str) and len(value) >= 2 and value[0] in {'"', "'"} and value[-1] == value[0]:
        try:
            return bytes(value[1:-1], "utf-8").decode("unicode_escape")
        except Exception:
            return value[1:-1]
    return value


class Parser:
    def __init__(self, source: str):
        self.source, self.tokens, self.i = source, lex(source), 0

    def peek(self) -> Token: return self.tokens[self.i]
    def pop(self) -> Token:
        token = self.tokens[self.i]; self.i += 1; return token
    def accept(self, value: str) -> bool:
        if self.peek().value == value: self.i += 1; return True
        return False

    def primary(self) -> Any:
        token = self.pop()
        if token.value == "(":
            value = self.expression(); self.accept(")"); return value
        if token.kind == "string": return _string(token.value)
        if token.kind == "number": return float(token.value)
        if token.kind == "ident":
            if self.accept("("):
                args = []
                while self.peek().kind != "eof" and self.peek().value != ")":
                    args.append(self.expression());
                    if not self.accept(","): break
                self.accept(")")
                return Call(token.value, args, token.line)
            return token.value
        if token.value == "[":
            values = []
            while self.peek().kind != "eof" and self.peek().value != "]":
                values.append(self.expression());
                if not self.accept(","): break
            self.accept("]"); return values
        if token.value == "{":
            obj: dict[str, Any] = {}
            while self.peek().kind != "eof" and self.peek().value != "}":
                key = _string(self.pop().value); self.accept(":"); obj[str(key)] = self.expression()
                if not self.accept(","): break
            self.accept("}"); return obj
        return token.value

    def expression(self, minimum: int = 0) -> Any:
        left = self.primary()
        precedence = {"+": 10, "-": 10, "*": 20, "/": 20}
        while self.peek().value in precedence and precedence[self.peek().value] >= minimum:
            op = self.pop(); right = self.expression(precedence[op.value] + 1)
            left = Call("__binary__", [left, op.value, right], op.line)
        return left

    def statements(self) -> list[Call]:
        found: list[Call] = []; scopes: list[dict[str, Any]] = [{}]

        def environment() -> dict[str, Any]:
            resolved: dict[str, Any] = {}
            for scope in scopes: resolved.update(scope)
            return resolved

        def assignment_scope(name: str) -> dict[str, Any]:
            return next((scope for scope in reversed(scopes) if name in scope), scopes[-1])

        while self.peek().kind != "eof":
            if self.accept("{"):
                scopes.append({}); continue
            if self.accept("}"):
                if len(scopes) > 1: scopes.pop()
                continue
            if self.peek().value == "var":
                self.pop()
                if self.peek().kind == "ident":
                    name = self.pop().value
                    if self.accept("="):
                        scopes[-1][name] = _resolve(self.expression(), environment())
                    else: scopes[-1][name] = name
                    if isinstance(scopes[-1][name], Call) and scopes[-1][name].name == "newSketch":
                        found.append(scopes[-1][name])
                continue
            if self.peek().kind == "ident" and self.i + 1 < len(self.tokens) and self.tokens[self.i + 1].value == "=":
                name = self.pop().value; self.pop()
                try:
                    scope = assignment_scope(name); scope[name] = _resolve(self.expression(), environment())
                    if isinstance(scope[name], Call) and scope[name].name == "newSketch":
                        value = scope[name]
                        found.append(value)
                except Exception: pass
            elif self.peek().kind == "ident" and self.i + 1 < len(self.tokens) and self.tokens[self.i + 1].value == "(":
                value = self.expression()
                if isinstance(value, Call):
                    value.args = [_resolve(arg, environment()) for arg in value.args]
                    found.append(value)
            else: self.i += 1
        return found


def _resolve(value: Any, environment: dict[str, Any], seen: set[str] | None = None) -> Any:
    seen = seen or set()
    if isinstance(value, str) and value in environment and value not in seen:
        return _resolve(environment[value], environment, seen | {value})
    if isinstance(value, Call): return Call(value.name, [_resolve(arg, environment, seen) for arg in value.args], value.line, value.raw)
    if isinstance(value, list): return [_resolve(arg, environment, seen) for arg in value]
    if isinstance(value, dict): return {key: _resolve(arg, environment, seen) for key, arg in value.items()}
    return value


def symbolic_string(value: Any) -> str:
    if isinstance(value, Call) and value.name == "__binary__" and value.args[1] == "+":
        return symbolic_string(value.args[0]) + symbolic_string(value.args[2])
    if isinstance(value, str): return "" if value == "id" else value
    return str(value)


def _arg_map(call: Call) -> dict[str, Any]:
    # Object literals are parsed as dictionaries; FeatureScript operation calls
    # conventionally put the definition map in the final argument.
    return next((arg for arg in reversed(call.args) if isinstance(arg, dict)), {})


def parse_featurescript(source: str, sample_id: str = "unknown") -> ModelIR:
    parser = Parser(source); calls = parser.statements()
    version = re.search(r"\bFeatureScript\s+(\d+(?:\.\d+)*)\s*;", source)
    standard_library = re.search(
        r"\bimport\s*\(\s*path\s*:\s*[\"']([^\"']*onshape/std/[^\"']*)[\"']",
        source,
    )
    model = ModelIR(
        sample_id,
        raw_source=source,
        featurescript_version=version.group(1) if version else None,
        standard_library=standard_library.group(1) if standard_library else None,
    )
    for call in calls:
        if call.name == "newSketch":
            definition = _arg_map(call)
            fid = symbolic_string(call.args[1]) if len(call.args) > 1 else f"sketch_{len(model.sketches)}"
            sketch_ir = SketchIR(fid, definition.get("sketchPlane"), [])
            model.sketches.append(sketch_ir); model.steps.append(sketch_ir)
        elif call.name in {"skLineSegment", "skCircle", "skArc", "skEllipse", "skFitSpline", "skPoint"}:
            # Attach sketch entities to the most recently declared sketch.
            if model.sketches:
                args = _arg_map(call); eid = str(call.args[1]) if len(call.args) > 1 else f"E{len(model.sketches[-1].entities)}"
                model.sketches[-1].entities.append(FeatureIR(eid, call.name, args, line_start=call.line, raw_source=call.name))
        elif call.name in {"extrude", "revolve", "fillet", "chamfer", "hole", "linearPattern", "mirror", "cPlane", "referenceAxis", "shell", "loft", "sweep", "circularPattern", "booleanBodies", "deleteBodies", "transform", "draft", "thicken", "split", "moveFace", "replaceFace", "deleteFace", "derive"}:
            fid = symbolic_string(call.args[1]) if len(call.args) > 1 else f"feature_{len(model.features)}"
            feature_ir = FeatureIR(fid, call.name, _arg_map(call), line_start=call.line, raw_source=call.name)
            model.features.append(feature_ir); model.steps.append(feature_ir)
    return model