#!/usr/bin/env python3
"""OpenOCD bridge utilities for EAB.

This is intentionally lightweight: it provides
- starting/stopping OpenOCD as a background process
- a tiny command channel via OpenOCD's telnet port (default 4444)

The goal is to keep all chip interaction "through EAB", including USB-JTAG workflows.
"""

from __future__ import annotations

import socket
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from .file_utils import write_json_file, tail_file
from .process_utils import pid_alive, read_pid_file, cleanup_pid_file, stop_process_graceful, popen_is_alive


DEFAULT_TELNET_PORT = 4444
DEFAULT_GDB_PORT = 3333
DEFAULT_TCL_PORT = 6666


def _scripts_dir() -> Optional[str]:
    # Homebrew default on Apple Silicon.
    p = Path("/opt/homebrew/share/openocd/scripts")
    if p.exists():
        return str(p)
    # Intel homebrew.
    p = Path("/usr/local/share/openocd/scripts")
    if p.exists():
        return str(p)
    return None


@dataclass(frozen=True)
class OpenOCDStatus:
    running: bool
    pid: Optional[int]
    cfg_path: Optional[str]
    log_path: str
    err_path: str
    last_error: Optional[str] = None
    telnet_port: int = DEFAULT_TELNET_PORT
    gdb_port: int = DEFAULT_GDB_PORT
    tcl_port: int = DEFAULT_TCL_PORT


class OpenOCDBridge:
    def __init__(self, base_dir: str):
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)
        self.pid_path = self.base_dir / "openocd.pid"
        self.status_path = self.base_dir / "openocd.status.json"
        self.log_path = self.base_dir / "openocd.log"
        self.err_path = self.base_dir / "openocd.err"
        self.cfg_path = self.base_dir / "openocd.cfg"

    def status(self) -> OpenOCDStatus:
        pid = read_pid_file(self.pid_path)
        running = bool(pid) and pid_alive(pid)
        if pid and not running:
            cleanup_pid_file(self.pid_path)
        return OpenOCDStatus(
            running=running,
            pid=pid if running else None,
            cfg_path=str(self.cfg_path) if self.cfg_path.exists() else None,
            log_path=str(self.log_path),
            err_path=str(self.err_path),
        )

    def _write_status(self, status: OpenOCDStatus) -> None:
        payload = {
            "running": status.running,
            "pid": status.pid,
            "cfg_path": status.cfg_path,
            "log_path": status.log_path,
            "err_path": status.err_path,
            "last_error": status.last_error,
            "telnet_port": status.telnet_port,
            "gdb_port": status.gdb_port,
            "tcl_port": status.tcl_port,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime()),
        }
        write_json_file(self.status_path, payload)

    def _ensure_cfg(
        self,
        *,
        chip: str,
        vid: str,
        pid: str,
        caps_descriptor: Optional[str],
        target_cfg: str,
        adapter: Optional[str] = None,
    ) -> str:
        """
        Generate OpenOCD config file for the specified chip.

        Supports:
        - ESP32 with USB-JTAG (esp_usb_jtag adapter)
        - STM32 with ST-Link (hla adapter)
        """
        chip_lower = chip.lower()

        lines = ["# Auto-generated by EAB (OpenOCDBridge)"]

        # Determine adapter type
        if adapter == "stlink" or chip_lower.startswith("stm32"):
            # STM32 with ST-Link
            lines.extend([
                "source [find interface/stlink.cfg]",
                "transport select hla_swd",
            ])
            # Determine target config based on STM32 variant
            if not target_cfg or target_cfg.startswith("target/esp"):
                stm32_target_map = {
                    "stm32f1": "target/stm32f1x.cfg",
                    "stm32f3": "target/stm32f3x.cfg",
                    "stm32f4": "target/stm32f4x.cfg",
                    "stm32l4": "target/stm32l4x.cfg",
                    "stm32h7": "target/stm32h7x.cfg",
                    "stm32g0": "target/stm32g0x.cfg",
                    "stm32g4": "target/stm32g4x.cfg",
                }
                for prefix, cfg in stm32_target_map.items():
                    if chip_lower.startswith(prefix):
                        target_cfg = cfg
                        break
                else:
                    target_cfg = "target/stm32l4x.cfg"  # Default
            lines.append(f"source [find {target_cfg}]")
            lines.append("reset_config srst_only")
        else:
            # ESP32 with USB-JTAG (default)
            lines.extend([
                f"set CHIPNAME {chip}",
                "adapter driver esp_usb_jtag",
                f"espusbjtag vid_pid {vid} {pid}",
            ])
            if caps_descriptor:
                lines.append(f"espusbjtag caps_descriptor {caps_descriptor}")
            lines.append(f"source [find {target_cfg}]")

        self.cfg_path.write_text("\n".join(lines) + "\n")
        return str(self.cfg_path)

    def start(
        self,
        *,
        chip: str = "esp32s3",
        # Espressif VID is typically 0x303a.
        vid: str = "0x303a",
        # Many ESP32-S3 "USB Serial/JTAG" devkits use 0x1001; bridge boards use 0x1002.
        pid: str = "0x1001",
        caps_descriptor: Optional[str] = "0x030A",
        target_cfg: str = "target/esp32s3.cfg",
        telnet_port: int = DEFAULT_TELNET_PORT,
        gdb_port: int = DEFAULT_GDB_PORT,
        tcl_port: int = DEFAULT_TCL_PORT,
        extra_args: Optional[list[str]] = None,
    ) -> OpenOCDStatus:
        cur = self.status()
        if cur.running:
            return cur

        scripts = _scripts_dir()
        cfg = self._ensure_cfg(
            chip=chip,
            vid=vid,
            pid=pid,
            caps_descriptor=caps_descriptor,
            target_cfg=target_cfg,
        )

        cmd = ["openocd"]
        if scripts:
            cmd += ["-s", scripts]
        cmd += ["-f", cfg]
        cmd += ["-c", f"telnet_port {telnet_port}"]
        cmd += ["-c", f"gdb_port {gdb_port}"]
        cmd += ["-c", f"tcl_port {tcl_port}"]
        if extra_args:
            cmd += extra_args

        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        log_f = self.log_path.open("a", encoding="utf-8")
        err_f = self.err_path.open("a", encoding="utf-8")

        proc = subprocess.Popen(
            cmd,
            stdout=log_f,
            stderr=err_f,
            cwd=str(self.base_dir),
            start_new_session=True,
        )
        self.pid_path.write_text(str(proc.pid))

        # Give it a moment to either bind ports or fail fast with a clear error.
        time.sleep(0.5)
        alive = pid_alive(proc.pid) and popen_is_alive(proc)
        last_error: Optional[str] = None
        if not alive:
            err_lines = tail_file(self.err_path, 20)
            last_error = "\n".join(err_lines).strip() or None
            cleanup_pid_file(self.pid_path)

        status = OpenOCDStatus(
            running=alive,
            pid=proc.pid if alive else None,
            cfg_path=cfg,
            log_path=str(self.log_path),
            err_path=str(self.err_path),
            last_error=last_error,
            telnet_port=telnet_port,
            gdb_port=gdb_port,
            tcl_port=tcl_port,
        )
        self._write_status(status)
        return status

    def stop(self, *, timeout_s: float = 5.0) -> OpenOCDStatus:
        cur = self.status()
        if not cur.running or not cur.pid:
            self._write_status(cur)
            cleanup_pid_file(self.pid_path)
            return cur

        stop_process_graceful(cur.pid, timeout_s)
        cleanup_pid_file(self.pid_path)

        status = OpenOCDStatus(
            running=False,
            pid=None,
            cfg_path=str(self.cfg_path) if self.cfg_path.exists() else None,
            log_path=str(self.log_path),
            err_path=str(self.err_path),
        )
        self._write_status(status)
        return status

    def cmd(self, command: str, *, timeout_s: float = 2.0, telnet_port: int = DEFAULT_TELNET_PORT) -> str:
        """Send a single command to OpenOCD via its telnet port.

        Returns the raw response (best-effort).
        """
        # The telnet server isn't a real telnet protocol implementation; it's a simple line-based
        # socket interface. We'll read until we see a prompt ('> ').
        with socket.create_connection(("127.0.0.1", telnet_port), timeout=timeout_s) as sock:
            sock.settimeout(timeout_s)
            # Drain banner/prompt.
            buf = b""
            start = time.time()
            while time.time() - start < timeout_s:
                try:
                    chunk = sock.recv(4096)
                except socket.timeout:
                    break
                if not chunk:
                    break
                buf += chunk
                if b"> " in buf:
                    break

            sock.sendall(command.encode("utf-8") + b"\n")
            out = b""
            start = time.time()
            while time.time() - start < timeout_s:
                chunk = sock.recv(4096)
                if not chunk:
                    break
                out += chunk
                if b"> " in out:
                    break

        # Strip trailing prompt.
        text = out.decode("utf-8", errors="replace")
        return text.replace("\r", "")
