from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from decimal import Decimal
from typing import Any, Dict, List, Optional


# -------------------- Lexer --------------------

class LexError(Exception):
    def __init__(self, message: str, line: int, col: int):
        super().__init__(f"LexError at {line}:{col} - {message}")
        self.line = line
        self.col = col


@dataclass(frozen=True)
class Token:
    kind: str
    value: Optional[str]
    line: int
    col: int


KEYWORDS = {"swap", "before", "after"}


def _is_ident_start(ch: str) -> bool:
    return ch.isalpha() or ch == "_"


def _is_ident_part(ch: str) -> bool:
    return ch.isalnum() or ch == "_"


def tokenise(source: str) -> List[Token]:
    tokens: List[Token] = []
    line = 1
    col = 1
    i = 0

    while i < len(source):
        ch = source[i]

        if ch == "\n":
            tokens.append(Token("NEWLINE", None, line, col))
            line += 1
            col = 1
            i += 1
            continue
        if ch == "\r":
            i += 1
            continue
        if ch.isspace():
            i += 1
            col += 1
            continue

        if ch == "#":
            while i < len(source) and source[i] not in ("\n", "\r"):
                i += 1
                col += 1
            continue

        if ch == "{":
            tokens.append(Token("LBRACE", "{", line, col))
            i += 1
            col += 1
            continue
        if ch == "}":
            tokens.append(Token("RBRACE", "}", line, col))
            i += 1
            col += 1
            continue
        if ch == ":":
            tokens.append(Token("COLON", ":", line, col))
            i += 1
            col += 1
            continue
        if ch == "=":
            tokens.append(Token("EQUAL", "=", line, col))
            i += 1
            col += 1
            continue

        if ch.isdigit():
            start_col = col
            start = i
            while i < len(source) and source[i].isdigit():
                i += 1
                col += 1
            if i < len(source) and source[i] == ".":
                i += 1
                col += 1
                if i >= len(source) or not source[i].isdigit():
                    raise LexError("Invalid number literal", line, col)
                while i < len(source) and source[i].isdigit():
                    i += 1
                    col += 1
            value = source[start:i]
            tokens.append(Token("NUMBER", value, line, start_col))
            continue

        if _is_ident_start(ch):
            start_col = col
            start = i
            i += 1
            col += 1
            while i < len(source) and _is_ident_part(source[i]):
                i += 1
                col += 1
            value = source[start:i]
            kind = "KEYWORD" if value in KEYWORDS else "IDENT"
            tokens.append(Token(kind, value, line, start_col))
            continue

        raise LexError(f"Unexpected character '{ch}'", line, col)

    tokens.append(Token("EOF", None, line, col))
    return tokens


# -------------------- Parser --------------------

class ParseError(Exception):
    def __init__(self, message: str, line: int, col: int):
        super().__init__(f"ParseError at {line}:{col} - {message}")
        self.line = line
        self.col = col


@dataclass(frozen=True)
class Entry:
    party: str
    amount: Decimal
    asset: str


@dataclass
class Program:
    before: List[Entry]
    after: List[Entry]


class Parser:
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0

    def _peek(self) -> Token:
        return self.tokens[self.pos]

    def _advance(self) -> Token:
        tok = self.tokens[self.pos]
        self.pos += 1
        return tok

    def _match(self, kind: str, value: Optional[str] = None) -> bool:
        tok = self._peek()
        if tok.kind != kind:
            return False
        if value is not None and tok.value != value:
            return False
        self._advance()
        return True

    def _expect(self, kind: str, value: Optional[str] = None) -> Token:
        tok = self._peek()
        if tok.kind != kind or (value is not None and tok.value != value):
            expected = f"{kind} {value}" if value else kind
            raise ParseError(f"Expected {expected}", tok.line, tok.col)
        return self._advance()

    def _skip_newlines(self) -> None:
        while self._match("NEWLINE"):
            pass

    def parse(self) -> Program:
        self._skip_newlines()
        self._expect("KEYWORD", "swap")
        self._skip_newlines()
        self._expect("LBRACE")

        before: List[Entry] = []
        after: List[Entry] = []
        seen_before = False
        seen_after = False

        while True:
            self._skip_newlines()
            if self._match("RBRACE"):
                break

            tok = self._peek()
            if tok.kind == "KEYWORD" and tok.value in ("before", "after"):
                section = tok.value
                if section == "before" and seen_before:
                    raise ParseError("Duplicate 'before' block", tok.line, tok.col)
                if section == "after" and seen_after:
                    raise ParseError("Duplicate 'after' block", tok.line, tok.col)

                self._advance()
                self._expect("COLON")
                if section == "before":
                    seen_before = True
                    before.extend(self._parse_entries_until_section_or_end())
                else:
                    seen_after = True
                    after.extend(self._parse_entries_until_section_or_end())
                continue

            raise ParseError("Expected 'before' or 'after' block", tok.line, tok.col)

        if not seen_before or not seen_after:
            tok = self._peek()
            raise ParseError("Both 'before' and 'after' blocks are required", tok.line, tok.col)

        return Program(before=before, after=after)

    def _parse_entries_until_section_or_end(self) -> List[Entry]:
        entries: List[Entry] = []
        while True:
            self._skip_newlines()
            tok = self._peek()
            if tok.kind == "RBRACE":
                break
            if tok.kind == "KEYWORD" and tok.value in ("before", "after"):
                break

            party_tok = self._expect("IDENT")
            self._expect("EQUAL")
            amount_tok = self._expect("NUMBER")
            asset_tok = self._expect("IDENT")
            entries.append(
                Entry(
                    party=party_tok.value,
                    amount=Decimal(amount_tok.value),
                    asset=asset_tok.value,
                )
            )
            self._skip_newlines()
        return entries


def parse(tokens: List[Token]) -> Program:
    return Parser(tokens).parse()


# -------------------- Verifier --------------------

class VerifyError(Exception):
    def __init__(self, message: str):
        super().__init__(f"VerifyError - {message}")


@dataclass
class VerificationSummary:
    valid: bool
    conservation: Dict[str, Dict[str, Decimal]]
    delta: Dict[str, Dict[str, Decimal]]


def verify(program: Program) -> VerificationSummary:
    _ensure_unique_parties(program.before, "before")
    _ensure_unique_parties(program.after, "after")

    for entry in program.before + program.after:
        if entry.amount < 0:
            raise VerifyError("Negative amounts are not allowed.")

    before_by_asset = _sum_by_asset(program.before)
    after_by_asset = _sum_by_asset(program.after)
    conservation = _build_conservation(before_by_asset, after_by_asset)

    return VerificationSummary(
        valid=True,
        conservation=conservation,
        delta=_compute_delta(program.before, program.after),
    )


def _sum_by_asset(entries: List[Entry]) -> Dict[str, Decimal]:
    sums: Dict[str, Decimal] = {}
    for e in entries:
        sums[e.asset] = sums.get(e.asset, Decimal("0")) + e.amount
    return sums


def _ensure_unique_parties(entries: List[Entry], label: str) -> None:
    seen = set()
    for e in entries:
        if e.party in seen:
            raise VerifyError(f"Duplicate party '{e.party}' in {label}.")
        seen.add(e.party)


def _build_conservation(
    before_by_asset: Dict[str, Decimal],
    after_by_asset: Dict[str, Decimal],
) -> Dict[str, Dict[str, Decimal]]:
    conservation: Dict[str, Dict[str, Decimal]] = {}
    assets = set(before_by_asset.keys()) | set(after_by_asset.keys())
    for asset in sorted(assets):
        before_val = before_by_asset.get(asset, Decimal("0"))
        after_val = after_by_asset.get(asset, Decimal("0"))
        if before_val != after_val:
            raise VerifyError(
                f"Conservation failed: {asset} before={_decimal_to_str(before_val)}, "
                f"after={_decimal_to_str(after_val)}"
            )
        conservation[asset] = {"before": before_val, "after": after_val}
    return conservation


def _compute_delta(before: List[Entry], after: List[Entry]) -> Dict[str, Dict[str, Decimal]]:
    delta: Dict[str, Dict[str, Decimal]] = {}
    for e in before:
        delta.setdefault(e.party, {})
        delta[e.party][e.asset] = delta[e.party].get(e.asset, Decimal("0")) - e.amount
    for e in after:
        delta.setdefault(e.party, {})
        delta[e.party][e.asset] = delta[e.party].get(e.asset, Decimal("0")) + e.amount
    return delta


# -------------------- Compiler --------------------

class CompileError(Exception):
    pass


@dataclass
class CompiledProof:
    program: Program
    summary: VerificationSummary

    def to_dict(self) -> Dict[str, Any]:
        return {
            "status": "valid" if self.summary.valid else "invalid",
            "conservation": _conservation_to_dict(self.summary.conservation),
            "delta": _delta_to_dict(self.summary.delta),
            "before": _entries_to_dict(self.program.before),
            "after": _entries_to_dict(self.program.after),
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), indent=2, sort_keys=True)


def compile_source(source: str) -> CompiledProof:
    try:
        tokens = tokenise(source)
        program = parse(tokens)
        summary = verify(program)
        return CompiledProof(program=program, summary=summary)
    except (LexError, ParseError, VerifyError) as exc:
        raise CompileError(str(exc)) from exc


def compile_file(path: str) -> CompiledProof:
    with open(path, "r", encoding="utf-8") as f:
        return compile_source(f.read())


def _entries_to_dict(entries: List[Entry]) -> Dict[str, str]:
    out: Dict[str, str] = {}
    for e in entries:
        out[e.party] = f"{_decimal_to_str(e.amount)} {e.asset}"
    return out


def _decimal_to_str(value: Decimal) -> str:
    text = format(value, "f")
    if "." in text:
        text = text.rstrip("0").rstrip(".")
    return text or "0"


def _decimal_to_signed_str(value: Decimal) -> str:
    base = _decimal_to_str(abs(value))
    sign = "+" if value >= 0 else "-"
    return f"{sign}{base}"


def _conservation_to_dict(conservation: Dict[str, Dict[str, Decimal]]) -> Dict[str, Dict[str, str]]:
    out: Dict[str, Dict[str, str]] = {}
    for asset, values in conservation.items():
        out[asset] = {
            "before": _decimal_to_str(values["before"]),
            "after": _decimal_to_str(values["after"]),
        }
    return out


def _delta_to_dict(delta: Dict[str, Dict[str, Decimal]]) -> Dict[str, Dict[str, str]]:
    out: Dict[str, Dict[str, str]] = {}
    for party, assets in delta.items():
        out[party] = {asset: _decimal_to_signed_str(amount) for asset, amount in assets.items()}
    return out


# Backwards-friendly alias
compile_swap = compile_source


# -------------------- CLI --------------------

def main() -> int:
    parser = argparse.ArgumentParser(
        prog="sophi",
        description="Sophi language compiler - generates portable validity proofs.",
    )
    sub = parser.add_subparsers(dest="command")

    compile_p = sub.add_parser("compile", help="Compile a .sph file into a proof")
    compile_p.add_argument("file", help="Path to .sph source file")
    compile_p.add_argument("-o", "--out", help="Write proof JSON to a file instead of stdout")

    args = parser.parse_args()

    if args.command != "compile":
        parser.print_help()
        return 2

    try:
        compiled = compile_file(args.file)
    except CompileError as exc:
        print(str(exc), file=sys.stderr)
        return 1

    output = compiled.to_json()
    if args.out:
        with open(args.out, "w", encoding="utf-8") as f:
            f.write(output)
    else:
        print(output)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
