from .tokens import TokenType


class ParseError(Exception):
    """Raised for syntax errors in the Prasad parser."""


class ASTNode:
    def __init__(self, line=None):
        self.line = line


class Program(ASTNode):
    def __init__(self, statements, line=None):
        super().__init__(line)
        self.statements = statements


class Assignment(ASTNode):
    def __init__(self, name, value, line=None):
        super().__init__(line)
        self.name = name
        self.value = value


class PrintStatement(ASTNode):
    def __init__(self, args, line=None):
        super().__init__(line)
        self.args = args


class IfStatement(ASTNode):
    def __init__(self, condition, body, else_body=None, line=None):
        super().__init__(line)
        self.condition = condition
        self.body = body
        self.else_body = else_body


class ForLoop(ASTNode):
    def __init__(self, condition, body, line=None):
        super().__init__(line)
        self.condition = condition
        self.body = body


class WhileStatement(ASTNode):
    def __init__(self, condition, body, line=None):
        super().__init__(line)
        self.condition = condition
        self.body = body


class FunctionDeclaration(ASTNode):
    def __init__(self, name, params, body, line=None):
        super().__init__(line)
        self.name = name
        self.params = params
        self.body = body


class ReturnStatement(ASTNode):
    def __init__(self, value, line=None):
        super().__init__(line)
        self.value = value


class BreakStatement(ASTNode):
    def __init__(self, line=None):
        super().__init__(line)


class ContinueStatement(ASTNode):
    def __init__(self, line=None):
        super().__init__(line)


class ThrowStatement(ASTNode):
    def __init__(self, value, line=None):
        super().__init__(line)
        self.value = value


class TryStatement(ASTNode):
    def __init__(
        self, try_body, catch_name=None, catch_body=None, finally_body=None, line=None
    ):
        super().__init__(line)
        self.try_body = try_body
        self.catch_name = catch_name
        self.catch_body = catch_body
        self.finally_body = finally_body


class FunctionCall(ASTNode):
    def __init__(self, name, args, line=None):
        super().__init__(line)
        self.name = name
        self.args = args


class Literal(ASTNode):
    def __init__(self, value, line=None):
        super().__init__(line)
        self.value = value


class FormattedString(ASTNode):
    def __init__(self, template, line=None):
        super().__init__(line)
        self.template = template


class Identifier(ASTNode):
    def __init__(self, name, line=None):
        super().__init__(line)
        self.name = name


class UnaryExpression(ASTNode):
    def __init__(self, op, expr, line=None):
        super().__init__(line)
        self.op = op
        self.expr = expr


class BinaryExpression(ASTNode):
    def __init__(self, left, op, right, line=None):
        super().__init__(line)
        self.left = left
        self.op = op
        self.right = right


class ListLiteral(ASTNode):
    def __init__(self, elements, line=None):
        super().__init__(line)
        self.elements = elements


class IndexExpression(ASTNode):
    def __init__(self, target, index, line=None):
        super().__init__(line)
        self.target = target
        self.index = index


class IndexAssignment(ASTNode):
    def __init__(self, target, index, value, line=None):
        super().__init__(line)
        self.target = target
        self.index = index
        self.value = value


class ObjectLiteral(ASTNode):
    def __init__(self, entries, line=None):
        super().__init__(line)
        self.entries = entries  # list[tuple[str, ASTNode]]


class Parser:
    def __init__(self, tokens, source=None):
        self.tokens = tokens
        self.position = 0
        self._source_lines = source.splitlines() if isinstance(source, str) else None

    def _token_str(self, token):
        """A human-friendly representation of a token for error messages."""
        if token.value is None:
            return token.type.name
        return str(token.value)

    def _raise_error(self, token, message, include_token=True):
        """Raise a ParseError with line/column context when source is available."""
        line = getattr(token, "line", None)
        col = getattr(token, "column", None)

        final_message = message
        if include_token:
            final_message = f"{final_message} '{self._token_str(token)}'"

        if (
            self._source_lines is not None
            and isinstance(line, int)
            and isinstance(col, int)
            and 1 <= line <= len(self._source_lines)
        ):
            src_line = self._source_lines[line - 1]
            caret = " " * max(col - 1, 0) + "^"
            raise ParseError(
                f"Error at line {line}, column {col}:\n{src_line}\n{caret}\n{final_message}"
            )

        if isinstance(line, int):
            raise ParseError(f"Chuk: {final_message}\nLine: {line}")
        raise ParseError(f"Chuk: {final_message}")

    def current(self):
        return self.tokens[self.position]

    def peek(self, offset=1):
        idx = self.position + offset
        if idx >= len(self.tokens):
            return self.tokens[-1]
        return self.tokens[idx]

    def advance(self):
        if self.position < len(self.tokens) - 1:
            self.position += 1
        return self.current()

    def match(self, type_, value=None):
        token = self.current()
        if token.type != type_:
            self._raise_error(token, f"Expected {type_.name}, got")
        if value is not None and token.value != value:
            self._raise_error(token, f"Expected {type_.name}('{value}'), got")
        self.advance()
        return token

    def match_name(self):
        token = self.current()
        if token.type not in (TokenType.IDENTIFIER, TokenType.KEYWORD):
            self._raise_error(token, "Expected identifier, got")
        self.advance()
        return token

    def parse(self):
        statements = []
        while self.current().type != TokenType.EOF:
            statements.append(self.statement())
        return Program(statements)

    def maybe_semicolon(self):
        if self.current().type == TokenType.SEMICOLON:
            self.advance()

    def statement(self):
        token = self.current()

        if token.type == TokenType.KEYWORD:
            if token.value == "dakhav":
                return self.print_statement()
            if token.value == "jar":
                return self.if_statement()
            if token.value == "sathi":
                return self.for_loop()
            if token.value == "joparyant":
                return self.while_statement()
            if token.value in ("kam", "karya"):
                return self.function_declaration()
            if token.value == "parat":
                return self.return_statement()
            if token.value == "todu":
                return self.break_statement()
            if token.value == "pudhe":
                return self.continue_statement()
            if token.value == "feka":
                return self.throw_statement()
            if token.value == "prayatna":
                return self.try_statement()

        if token.type == TokenType.IDENTIFIER:
            if self.peek().type == TokenType.OPERATOR and self.peek().value == "=":
                return self.assignment()

        if token.type in (TokenType.IDENTIFIER, TokenType.KEYWORD):
            if self._looks_like_call():
                return self.function_call()
            if token.type == TokenType.IDENTIFIER and self.peek().type == TokenType.LBRACKET:
                expr = self.assignment_expression()
                if isinstance(expr, IndexAssignment):
                    self.maybe_semicolon()
                    return expr

        self._raise_error(token, "Unexpected token")

    def assignment(self):
        name_token = self.match(TokenType.IDENTIFIER)
        name = name_token.value
        self.match(TokenType.OPERATOR, "=")
        value = self.expression()
        self.maybe_semicolon()
        return Assignment(name, value, line=name_token.line)

    def print_statement(self):
        kw = self.match(TokenType.KEYWORD, "dakhav")
        self.match(TokenType.LPAREN)

        args = []
        if self.current().type != TokenType.RPAREN:
            args.append(self.expression())
            while self.current().type == TokenType.COMMA:
                self.advance()
                args.append(self.expression())

        self.match(TokenType.RPAREN)
        self.maybe_semicolon()
        return PrintStatement(args, line=kw.line)

    def _looks_like_call(self):
        idx = self.position
        if idx >= len(self.tokens):
            return False
        if self.tokens[idx].type not in (TokenType.IDENTIFIER, TokenType.KEYWORD):
            return False
        idx += 1

        while idx < len(self.tokens) and self.tokens[idx].type == TokenType.DOT:
            idx += 1
            if idx >= len(self.tokens):
                return False
            if self.tokens[idx].type not in (TokenType.IDENTIFIER, TokenType.KEYWORD):
                return False
            idx += 1

        return idx < len(self.tokens) and self.tokens[idx].type == TokenType.LPAREN

    def _parse_call(self, *, allow_semicolon: bool):
        name_token = self.match_name()
        parts = [name_token.value]

        while self.current().type == TokenType.DOT:
            self.advance()
            parts.append(self.match_name().value)

        name = ".".join(parts)
        self.match(TokenType.LPAREN)

        args = []
        if self.current().type != TokenType.RPAREN:
            args.append(self.expression())
            while self.current().type == TokenType.COMMA:
                self.advance()
                args.append(self.expression())

        self.match(TokenType.RPAREN)
        if allow_semicolon:
            self.maybe_semicolon()
        return FunctionCall(name, args, line=name_token.line)

    def function_call(self):
        return self._parse_call(allow_semicolon=True)

    def if_statement(self):
        self.match(TokenType.KEYWORD, "jar")
        condition = self.expression()
        body = self.block()

        else_body = None

        # Support else-if chains:
        # jar (...) { ... }
        # nahi_tar_jar (...) { ... }
        # nahitar_jar (...) { ... }
        # nahitar { ... }
        tail_if = None
        while (
            self.current().type == TokenType.KEYWORD
            and self.current().value in ("nahi_tar_jar", "nahitar_jar")
        ):
            kw = self.current()
            self.advance()
            elif_condition = self.expression()
            elif_body = self.block()
            elif_node = IfStatement(elif_condition, elif_body, else_body=None, line=kw.line)

            if else_body is None:
                else_body = [elif_node]
            else:
                tail_if.else_body = [elif_node]

            tail_if = elif_node

        if self.current().type == TokenType.KEYWORD and self.current().value == "nahitar":
            self.advance()
            block = self.block()
            if tail_if is None:
                else_body = block
            else:
                tail_if.else_body = block

        return IfStatement(condition, body, else_body)

    def for_loop(self):
        # Placeholder grammar; kept for forward-compat.
        self.match(TokenType.KEYWORD, "sathi")
        condition = self.expression()
        body = self.block()
        return ForLoop(condition, body)

    def while_statement(self):
        self.match(TokenType.KEYWORD, "joparyant")
        condition = self.expression()
        body = self.block()
        return WhileStatement(condition, body)

    def function_declaration(self):
        # Placeholder grammar; kept for forward-compat.
        if self.current().type != TokenType.KEYWORD or self.current().value not in (
            "kam",
            "karya",
        ):
            self._raise_error(self.current(), "Expected function keyword, got")
        self.advance()
        name = self.match_name().value
        self.match(TokenType.LPAREN)

        params = []
        if self.current().type != TokenType.RPAREN:
            params.append(self.match_name().value)
            while self.current().type == TokenType.COMMA:
                self.advance()
                params.append(self.match_name().value)

        self.match(TokenType.RPAREN)
        body = self.block()
        return FunctionDeclaration(name, params, body)

    def return_statement(self):
        self.match(TokenType.KEYWORD, "parat")
        value = self.expression()
        self.maybe_semicolon()
        return ReturnStatement(value)

    def break_statement(self):
        kw = self.match(TokenType.KEYWORD, "todu")
        self.maybe_semicolon()
        return BreakStatement(line=kw.line)

    def continue_statement(self):
        kw = self.match(TokenType.KEYWORD, "pudhe")
        self.maybe_semicolon()
        return ContinueStatement(line=kw.line)

    def throw_statement(self):
        kw = self.match(TokenType.KEYWORD, "feka")
        value = self.expression()
        self.maybe_semicolon()
        return ThrowStatement(value, line=kw.line)

    def try_statement(self):
        kw = self.match(TokenType.KEYWORD, "prayatna")
        try_body = self.block()

        catch_name = None
        catch_body = None
        finally_body = None

        if self.current().type == TokenType.KEYWORD and self.current().value == "pakad":
            self.advance()
            if self.current().type == TokenType.LPAREN:
                self.advance()
                catch_name = self.match(TokenType.IDENTIFIER).value
                self.match(TokenType.RPAREN)
            catch_body = self.block()

        if self.current().type == TokenType.KEYWORD and self.current().value == "shevati":
            self.advance()
            finally_body = self.block()

        if catch_body is None and finally_body is None:
            self._raise_error(self.current(), "Expected pakad or shevati after prayatna")

        return TryStatement(
            try_body,
            catch_name=catch_name,
            catch_body=catch_body,
            finally_body=finally_body,
            line=kw.line,
        )

    def block(self):
        self.match(TokenType.LBRACE)
        statements = []
        while self.current().type != TokenType.RBRACE:
            statements.append(self.statement())
        self.match(TokenType.RBRACE)
        return statements

    def expression(self):
        return self.assignment_expression()

    def assignment_expression(self):
        expr = self.logical_or()

        if self.current().type == TokenType.OPERATOR and self.current().value == "=":
            token = self.current()
            if isinstance(expr, Identifier):
                self.advance()
                value = self.assignment_expression()
                return Assignment(expr.name, value, line=token.line)
            if isinstance(expr, IndexExpression):
                self.advance()
                value = self.assignment_expression()
                return IndexAssignment(expr.target, expr.index, value, line=token.line)

            self._raise_error(token, "Invalid assignment target")

        return expr

    def logical_or(self):
        expr = self.logical_and()
        while (
            self.current().type == TokenType.KEYWORD and self.current().value == "kiwa"
        ):
            token = self.current()
            op = "||"
            self.advance()
            right = self.logical_and()
            expr = BinaryExpression(expr, op, right, line=token.line)
        return expr

    def logical_and(self):
        expr = self.equality()
        while (
            self.current().type == TokenType.KEYWORD and self.current().value == "aani"
        ):
            token = self.current()
            op = "&&"
            self.advance()
            right = self.equality()
            expr = BinaryExpression(expr, op, right, line=token.line)
        return expr

    def equality(self):
        expr = self.comparison()
        while self.current().type == TokenType.OPERATOR and self.current().value in (
            "==",
            "!=",
        ):
            token = self.current()
            op = token.value
            self.advance()
            right = self.comparison()
            expr = BinaryExpression(expr, op, right, line=token.line)
        return expr

    def comparison(self):
        expr = self.term()
        while self.current().type == TokenType.OPERATOR and self.current().value in (
            ">",
            "<",
            ">=",
            "<=",
        ):
            token = self.current()
            op = token.value
            self.advance()
            right = self.term()
            expr = BinaryExpression(expr, op, right, line=token.line)
        return expr

    def term(self):
        expr = self.factor()
        while self.current().type == TokenType.OPERATOR and self.current().value in (
            "+",
            "-",
        ):
            token = self.current()
            op = token.value
            self.advance()
            right = self.factor()
            expr = BinaryExpression(expr, op, right, line=token.line)
        return expr

    def factor(self):
        expr = self.unary()
        while self.current().type == TokenType.OPERATOR and self.current().value in (
            "*",
            "/",
            "%",
        ):
            token = self.current()
            op = token.value
            self.advance()
            right = self.unary()
            expr = BinaryExpression(expr, op, right, line=token.line)
        return expr

    def unary(self):
        if self.current().type == TokenType.OPERATOR and self.current().value in (
            "-",
            "+",
            "!",
        ):
            token = self.current()
            op = token.value
            self.advance()
            return UnaryExpression(op, self.unary(), line=token.line)
        if self.current().type == TokenType.KEYWORD and self.current().value == "navhe":
            token = self.current()
            self.advance()
            return UnaryExpression("!", self.unary(), line=token.line)
        return self.primary()

    def primary(self):
        token = self.current()

        if (
            token.type == TokenType.IDENTIFIER
            and token.value == "f"
            and self.peek().type == TokenType.STRING
        ):
            f_token = self.match(TokenType.IDENTIFIER)
            template_token = self.match(TokenType.STRING)
            expr = FormattedString(template_token.value, line=f_token.line)
        elif token.type == TokenType.KEYWORD and token.value in ("ho", "khara"):
            self.advance()
            expr = Literal(True, line=token.line)
        elif token.type == TokenType.KEYWORD and token.value in ("nahi", "khota"):
            self.advance()
            expr = Literal(False, line=token.line)
        elif token.type == TokenType.KEYWORD and token.value == "kahi":
            self.advance()
            expr = Literal(None, line=token.line)
        elif token.type == TokenType.NUMBER:
            self.advance()
            expr = Literal(token.value, line=token.line)
        elif token.type == TokenType.STRING:
            self.advance()
            expr = Literal(token.value, line=token.line)
        elif token.type in (TokenType.IDENTIFIER, TokenType.KEYWORD):
            if self._looks_like_call():
                # Expression-level function calls (no trailing semicolon here).
                expr = self._parse_call(allow_semicolon=False)
            else:
                self.advance()
                expr = Identifier(token.value, line=token.line)
        elif token.type == TokenType.LPAREN:
            self.advance()
            expr = self.expression()
            self.match(TokenType.RPAREN)
        elif token.type == TokenType.LBRACKET:
            expr = self.list_literal()
        elif token.type == TokenType.LBRACE:
            expr = self.object_literal()
        else:
            if token.type == TokenType.EOF:
                self._raise_error(token, "Expected expression", include_token=False)
            self._raise_error(token, "Expected expression, got")

        while self.current().type == TokenType.LBRACKET:
            bracket = self.match(TokenType.LBRACKET)
            index = self.expression()
            self.match(TokenType.RBRACKET)
            expr = IndexExpression(expr, index, line=bracket.line)

        return expr

    def list_literal(self):
        start = self.match(TokenType.LBRACKET)
        elements = []
        if self.current().type != TokenType.RBRACKET:
            elements.append(self.expression())
            while self.current().type == TokenType.COMMA:
                self.advance()
                elements.append(self.expression())
        self.match(TokenType.RBRACKET)
        return ListLiteral(elements, line=start.line)

    def object_literal(self):
        start = self.match(TokenType.LBRACE)
        entries = []

        if self.current().type != TokenType.RBRACE:
            while True:
                key_token = self.current()
                if key_token.type != TokenType.STRING:
                    self._raise_error(key_token, "Object keys must be strings, got")
                self.advance()

                self.match(TokenType.COLON)
                value = self.expression()
                entries.append((key_token.value, value))

                if self.current().type != TokenType.COMMA:
                    break
                self.advance()

        self.match(TokenType.RBRACE)
        return ObjectLiteral(entries, line=start.line)
