import os
import sys
import ast
import json
import re
import hashlib
import time
import uuid
import tempfile
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, Dict, Any, List
from dataclasses import dataclass, field, asdict

CODE_VERSION = "0.2.0-enterprise"
SCHEMA_VERSION = "2.0.0"
DETERMINISTIC_MODE = os.getenv("DETERMINISTIC_MODE", "true").lower() in ("1", "true", "yes")
CHECKPOINT_DIR = Path(os.getenv("OSTS_CHECKPOINT_DIR", "./checkpoints"))
CACHE_DIR = Path(os.getenv("OSTS_CACHE_DIR", "./translation_cache"))

_TEMPLATE = """/**
 * @generated-by    : osts-translator
 * @trace_id        : {trace_id}
 * @schema_version  : {schema_version}
 * @code_version    : {code_version}
 * @deterministic   : {deterministic}
 * @src_hash        : {src_hash}
 * @direction       : {direction}
 * @timestamp       : {timestamp}
 */
"""


@dataclass(frozen=True)
class IROperation:
    op_type: str
    properties: Dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class IRNode:
    lang: str
    trace_id: str
    source_hash: str
    schema_version: str = SCHEMA_VERSION
    nodes: List[IROperation] = field(default_factory=list)
    meta: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict:
        return {
            "lang": self.lang,
            "trace_id": self.trace_id,
            "source_hash": self.source_hash,
            "schema_version": self.schema_version,
            "nodes": [asdict(n) for n in self.nodes],
            "meta": self.meta,
        }

    @classmethod
    def from_dict(cls, d: Dict) -> "IRNode":
        nodes = [IROperation(**n) for n in d.get("nodes", [])]
        return cls(
            lang=d["lang"],
            trace_id=d["trace_id"],
            source_hash=d["source_hash"],
            schema_version=d.get("schema_version", SCHEMA_VERSION),
            nodes=nodes,
            meta=d.get("meta", {}),
        )


def _now_iso() -> str:
    return (
        datetime.now(timezone.utc).isoformat()
        if not DETERMINISTIC_MODE
        else "1970-01-01T00:00:00+00:00"
    )


def _trace_id(source: str, direction: str) -> str:
    seed = f"{SCHEMA_VERSION}:{direction}:{source}"
    if DETERMINISTIC_MODE:
        return hashlib.sha256(seed.encode()).hexdigest()[:32]
    return str(uuid.uuid4())


def _hash_source(source: str) -> str:
    return hashlib.sha256(source.encode()).hexdigest()


def atomic_write(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(content, encoding="utf-8")
    os.replace(str(tmp), str(path))


def structured_log(trace_id: str, stage: str, status: str, details: Optional[Dict] = None):
    entry = {
        "ts": _now_iso(),
        "trace_id": trace_id,
        "code_version": CODE_VERSION,
        "schema_version": SCHEMA_VERSION,
        "stage": stage,
        "status": status,
        "details": details or {},
    }
    print(json.dumps(entry, default=str), flush=True)


def checkpoint_path(trace_id: str, stage: str) -> Path:
    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    return CHECKPOINT_DIR / f"{trace_id}.{stage}.json"


def save_checkpoint(trace_id: str, stage: str, data: Dict):
    atomic_write(checkpoint_path(trace_id, stage), json.dumps(data, indent=2))


def load_checkpoint(trace_id: str, stage: str) -> Optional[Dict]:
    cp = checkpoint_path(trace_id, stage)
    if cp.exists():
        return json.loads(cp.read_text(encoding="utf-8"))
    return None


def cache_path(trace_id: str) -> Path:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    return CACHE_DIR / f"{trace_id}.translated.txt"


class PyAstWalker(ast.NodeVisitor):
    def __init__(self, source: str):
        self._source = source
        self._tree: ast.AST = ast.parse(source)
        self._nodes: List[Dict] = []

    def build_ir_ops(self) -> List[IROperation]:
        self.visit(self._tree)
        return [IROperation(**n) for n in self._nodes]

    def visit_Import(self, node):
        for alias in node.names:
            self._nodes.append({"op_type": "import", "properties": {"name": alias.name}})
        self.generic_visit(node)

    def visit_ImportFrom(self, node):
        if node.module:
            for alias in node.names:
                self._nodes.append(
                    {"op_type": "import", "properties": {"name": f"{node.module}.{alias.name}"}}
                )
        self.generic_visit(node)

    def visit_Call(self, node):
        if isinstance(node.func, ast.Attribute):
            chain = _attr_chain(node.func)
            if any(p in chain for p in ("read_excel", "to_excel")):
                kwargs = {}
                for kw in node.keywords:
                    try:
                        kwargs[kw.arg] = ast.unparse(kw.value)
                    except Exception:
                        kwargs[kw.arg] = str(kw.value)
                self._nodes.append(
                    {"op_type": "excel_io", "properties": {"pandas_method": chain, "kwargs": kwargs}}
                )
            elif "load_workbook" in chain:
                self._nodes.append({"op_type": "openpyxl_load", "properties": {"method": chain}})
            elif "save" in chain:
                self._nodes.append({"op_type": "openpyxl_save", "properties": {"method": chain}})
        self.generic_visit(node)


def _attr_chain(node: ast.Attribute) -> str:
    parts = []
    n = node
    while isinstance(n, ast.Attribute):
        parts.append(n.attr)
        n = n.value
    if isinstance(n, ast.Name):
        parts.append(n.id)
    return ".".join(reversed(parts))


def parse_python(source: str, trace_id: str) -> IRNode:
    ops = PyAstWalker(source).build_ir_ops()
    return IRNode(lang="python", trace_id=trace_id, source_hash=_hash_source(source), nodes=ops)


class OSTSRegexExtractor:
    _RE_WS = re.compile(
        r'workbook\.(?:getWorksheet|addWorksheet)\s*\(\s*["\'](.+?)["\']\s*\)'
    )
    _RE_LET = re.compile(r'let\s+(\w+)\s*=\s*(.+?);')
    _RE_RANGE = re.compile(r'\.(getUsedRange|getValues|setValues)\s*\(\s*\)')

    def __init__(self, source: str):
        self._src = source

    def to_ir(self, trace_id: str) -> IRNode:
        nodes: List[IROperation] = []
        for m in self._RE_WS.finditer(self._src):
            nodes.append(
                IROperation(
                    "excel_io",
                    {"action": "get_or_add_worksheet", "sheet_expr": m.group(1)},
                )
            )
        for m in self._RE_LET.finditer(self._src):
            nodes.append(IROperation("declaration", {"var": m.group(1), "expr": m.group(2)}))
        for m in self._RE_RANGE.finditer(self._src):
            nodes.append(IROperation("range_op", {"op": m.group(1)}))
        return IRNode(
            lang="osts",
            trace_id=trace_id,
            source_hash=_hash_source(self._src),
            nodes=nodes,
        )


def parse_osts(source: str, trace_id: str) -> IRNode:
    return OSTSRegexExtractor(source).to_ir(trace_id)


def _header(trace_id: str, src_hash: str, direction: str) -> str:
    return _TEMPLATE.format(
        trace_id=trace_id,
        schema_version=SCHEMA_VERSION,
        code_version=CODE_VERSION,
        deterministic=str(DETERMINISTIC_MODE).lower(),
        src_hash=src_hash,
        direction=direction,
        timestamp=_now_iso(),
    )


def emit_osts(ir: IRNode) -> str:
    lines = [
        _header(ir.trace_id, ir.source_hash, "ir->osts"),
        "function main(workbook: ExcelScript.Workbook) {",
    ]
    for op in ir.nodes:
        t = op.op_type
        props = op.properties
        if t == "import":
            lines.append(f"  // enterprise: dep '{props.get('name')}'")
        elif t == "excel_io":
            method = props.get("pandas_method", "")
            kwargs = props.get("kwargs", {})
            sheet = kwargs.get("sheet_name", "'Sheet1'")
            if "read_excel" in method:
                lines.append(f"  let ws = workbook.getWorksheet({sheet});")
                lines.append("  let table = ws.addTable(ws.getUsedRange(), true);")
            elif "to_excel" in method:
                lines.append(f"  let outWs = workbook.getWorksheet({sheet}) || workbook.addWorksheet({sheet});")
                lines.append("  // [TRANSLATION NOTE] map DataFrame rows to range.setValues()")
            else:
                lines.append(f"  // excel_io: {method}")
        elif t == "openpyxl_load":
            lines.append("  // openpyxl.load_workbook maps to incoming workbook arg")
        elif t == "openpyxl_save":
            lines.append("  // openpyxl.save maps to implicit workbook save context")
        elif t == "range_op":
            op_name = props.get("op", "rangeOp")
            lines.append(f"  // [REVIEW] OSTS range op '{op_name}' requires manual review")
        else:
            lines.append(f"  // UNMAPPED: {t} {json.dumps(props)}")
    lines.append("}")
    return "\n".join(lines)


def emit_python(ir: IRNode) -> str:
    lines = [
        _header(ir.trace_id, ir.source_hash, "osts->python"),
        "import openpyxl",
        "from openpyxl import Workbook",
        "",
        "def main(file_path: str):",
        '    wb = openpyxl.load_workbook(file_path)',
    ]
    for op in ir.nodes:
        t = op.op_type
        props = op.properties
        if t == "excel_io" and props.get("action") == "get_or_add_worksheet":
            sheet = props.get("sheet_expr", "Sheet1")
            lines.append(f'    ws = wb["{sheet}"]')
        elif t == "declaration":
            lines.append(f"    # let {props['var']} = {props['expr']}  # review needed")
        elif t == "range_op":
            op_name = props.get("op", "")
            if op_name == "getUsedRange":
                lines.append("    used = ws.dimensions")
            elif op_name == "getValues":
                lines.append("    data = [ [cell.value for cell in row] for row in ws.iter_rows() ]")
            else:
                lines.append(f"    # range op '{op_name}'")
        else:
            lines.append(f"    # node: {t}")
    lines.extend([
        '    wb.save(file_path)',
        "",
        'if __name__ == "__main__":',
        '    main("example.xlsx")',
    ])
    return "\n".join(lines)


@dataclass(frozen=True)
class DirectionSpec:
    parser: Any
    emitter: Any


DIRECTION_REGISTRY = {
    "Python → OSTS": DirectionSpec(parser=parse_python, emitter=emit_osts),
    "OSTS → Python": DirectionSpec(parser=parse_osts, emitter=emit_python),
}


class TranslationPipeline:
    def __init__(self, source: str, direction: str):
        self.source = source
        self.direction = direction
        self.trace_id = _trace_id(source, direction)
        self.spec = DIRECTION_REGISTRY[direction]
        self._src_hash = _hash_source(source)

    def run(self) -> str:
        cpath = cache_path(self.trace_id)
        if DETERMINISTIC_MODE and cpath.exists():
            structured_log(self.trace_id, "cache", "hit")
            return cpath.read_text(encoding="utf-8")

        stage = "parse"
        cp = load_checkpoint(self.trace_id, stage)
        if cp:
            ir = IRNode.from_dict(cp)
            structured_log(self.trace_id, stage, "resumed_from_checkpoint")
        else:
            ir = self.spec.parser(self.source, self.trace_id)
            save_checkpoint(self.trace_id, stage, ir.to_dict())
            structured_log(self.trace_id, stage, "completed", {"nodes": len(ir.nodes)})

        stage = "emit"
        cp = load_checkpoint(self.trace_id, stage)
        if cp and cp.get("output"):
            output = cp["output"]
            structured_log(self.trace_id, stage, "resumed_from_checkpoint")
        else:
            output = self.spec.emitter(ir)
            save_checkpoint(self.trace_id, stage, {"output": output, "version": CODE_VERSION})
            structured_log(self.trace_id, stage, "completed", {"chars": len(output)})

        if DETERMINISTIC_MODE:
            atomic_write(cpath, output)
            structured_log(self.trace_id, "cache", "write", {"path": str(cpath)})

        return output


import gradio as gr

PY_SAMPLE = """import pandas as pd
df = pd.read_excel("data.xlsx", sheet_name="Sales")
df_filtered = df[df["Amount"] > 1000]
df_filtered.to_excel("output.xlsx", sheet_name="Filtered", index=False)
"""

OSTS_SAMPLE = """function main(workbook: ExcelScript.Workbook) {
  let ws = workbook.getWorksheet("Sales");
  let table = ws.addTable(ws.getUsedRange(), true);
  let newWs = workbook.addWorksheet("Filtered");
  // filter and write back
}
"""


def translate_ui(code: str, direction: str) -> str:
    if not code.strip():
        return "ERROR: Empty input."
    pipeline = TranslationPipeline(code, direction)
    try:
        return pipeline.run()
    except Exception as e:
        import traceback
        structured_log(pipeline.trace_id, "orchestrator", "fatal")
        return f"Pipeline error:\n{traceback.format_exc()}"


if __name__ == "__main__":
    with gr.Blocks(title="OSTS Enterprise Translator") as demo:
        gr.Markdown("## OSTS Enterprise Translator")
        gr.Markdown(
            "Behavioral tags: deterministic, idempotent, schema-bound, checkpointed, "
            "traceable, versioned, atomic, stateless, pure, modular, composable, "
            "encapsulated, declarative, transactional, retry-safe, predictable."
        )
        with gr.Row():
            with gr.Column(scale=1):
                direction = gr.Dropdown(
                    choices=list(DIRECTION_REGISTRY.keys()),
                    value="Python → OSTS",
                    label="Direction",
                )
                code_in = gr.Code(label="Input", language="python", value=PY_SAMPLE)
                btn = gr.Button("Translate", variant="primary")

            with gr.Column(scale=1):
                code_out = gr.Code(label="Output", language="typescript")

        btn.click(fn=translate_ui, inputs=[code_in, direction], outputs=code_out)

        gr.Examples(
            examples=[
                [PY_SAMPLE, "Python → OSTS"],
                [OSTS_SAMPLE, "OSTS → Python"],
            ],
            inputs=[code_in, direction],
            label="Quick Samples",
        )

    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )