#!/usr/bin/env python3

import os
import sys
import io
import time
import threading
import cv2
import numpy as np
import subprocess
import re
import tempfile
import requests
from pathlib import Path
from datetime import datetime
import yaml
import signal

try:
    import PIL.Image
    HAS_PIL = True
except ImportError:
    HAS_PIL = False

API_URL = os.environ.get("API_URL", "https://simlink.sirblob.co")

class RunRecorder:
    def __init__(self, results_dir=None, fps=5):
        project_dir = Path(__file__).resolve().parent.parent.parent
        self.fps = fps
        self.start_time = time.time()

        self.search_start = 0.0
        self.search_duration = 0.0

        config_data = {}
        for cfg in ["search.yaml", "ugv.yaml", "uav.yaml"]:
            p = project_dir / "config" / cfg
            if p.exists():
                try:
                    with open(p, 'r') as f:
                        data = yaml.safe_load(f)
                        if data:
                            key = cfg.replace('.yaml', '')
                            config_data[key] = data
                except Exception as e:
                    print(f"[REC] Failed to load config {cfg}: {e}")

        if "ugv" in config_data and "topics" in config_data["ugv"]:
            del config_data["ugv"]["topics"]

        if "uav" in config_data and "connection" in config_data["uav"]:
            del config_data["uav"]["connection"]

        self.sim_id = 0
        self.sim_name = "simulation_unknown"
        try:
            resp = requests.post(f"{API_URL}/api/simulations/create", json=config_data, timeout=10)
            if resp.status_code == 200:
                data = resp.json()
                self.sim_id = data.get("id", 0)
                self.sim_name = f"simulation_{self.sim_id}"
        except Exception as e:
            print(f"[REC] API create failed: {e}")

        if results_dir:
            self.run_dir = Path(results_dir) / self.sim_name
        else:
            self.run_dir = project_dir / "results" / self.sim_name
        self.run_dir.mkdir(parents=True, exist_ok=True)

        self.run_num = self.sim_id

        print(f"[REC] Recording local results to {self.run_dir.name} and API ({self.sim_name}, ID: {self.sim_id})")
        self._log_path = self.run_dir / "log.txt"
        self._log_file = open(self._log_path, "w")
        self._original_stdout = sys.stdout
        self._original_stderr = sys.stderr
        self._tee_stdout = _TeeWriter(sys.stdout, self._log_file)
        self._tee_stderr = _TeeWriter(sys.stderr, self._log_file)

        self._tracker_writer = None
        self._camera_writer = None
        self._ugv_camera_writer = None
        self._gazebo_writer = None

        self._tracker_size = None
        self._camera_size = None
        self._ugv_camera_size = None
        self._gazebo_size = None

        self._tracker_frames = 0
        self._camera_frames = 0
        self._ugv_camera_frames = 0
        self._gazebo_frames = 0

        self._last_tracker_frame = None
        self._last_camera_frame = None
        self._last_ugv_camera_frame = None
        self._camera_snapshots = []

        self._recording = False
        self._tracker_ref = None
        self._camera_ref = None
        self._record_thread = None
        self._lock = threading.Lock()

        threading.Thread(target=self._upload_hardware_info, daemon=True).start()

    def _upload_hardware_info(self):
        if not self.sim_id:
            return

        payload = {
            "cpu_info": "Unknown CPU",
            "gpu_info": "Unknown GPU",
            "ram_info": "Unknown RAM"
        }

        try:
            import subprocess
            cpu = "Unknown"
            try:
                cpu = subprocess.check_output("grep -m 1 'model name' /proc/cpuinfo | cut -d ':' -f 2", shell=True, timeout=2).decode('utf-8').strip()
            except Exception: pass
            if cpu: payload["cpu_info"] = cpu

            try:
                ram_kb = int(subprocess.check_output("awk '/MemTotal/ {print $2}' /proc/meminfo", shell=True, timeout=2))
                payload["ram_info"] = f"{round(ram_kb / 1024 / 1024, 1)} GB"
            except Exception: pass

            try:
                gpu = subprocess.check_output("lspci | grep -i vga | cut -d ':' -f 3", shell=True, timeout=2).decode('utf-8').strip()
                if gpu: payload["gpu_info"] = gpu
            except Exception: pass
        except Exception:
            pass

        try:
            requests.put(f"{API_URL}/api/simulations/{self.sim_id}/hardware", json=payload, timeout=5)
        except Exception as e:
            print(f"[REC] Hardware info sync failed: {e}")

    def _upload_file(self, path, filename):
        if not self.sim_id:
            return
        try:
            with open(path, 'rb') as f:
                resp = requests.post(
                    f"{API_URL}/api/simulations/{self.sim_id}/upload",
                    files={"file": (filename, f)},
                    timeout=60
                )
                if resp.status_code == 200:
                    pass
        except Exception as e:
            print(f"[REC] Upload failed for {filename}: {e}")

    def _find_gazebo_window_ffmpeg(self):
        try:
            cmd = "xwininfo -root -tree | grep -i 'gazebo\|gz sim'"
            output = subprocess.check_output(cmd, shell=True).decode('utf-8')
            lines = [l for l in output.strip().split('\n') if "1x1" not in l and "Gazebo Sim" in l]
            if not lines:
                return None

            wid_line = lines[0]
            wid_match = re.search(r'(0x[0-9a-fA-F]+)', wid_line)
            if wid_match:
                return wid_match.group(1)
        except Exception:
            pass
        return None

    def _get_window_geometry(self, window_id):
        try:
            geo = subprocess.run(
                ["xdotool", "getwindowgeometry", window_id],
                capture_output=True, text=True
            )
            x = y = width = height = None
            for line in geo.stdout.splitlines():
                if "Position:" in line:
                    pos = line.split("Position:")[1].strip().split()[0]
                    x, y = map(int, pos.split(","))
                elif "Geometry:" in line:
                    size = line.split("Geometry:")[1].strip()
                    width, height = map(int, size.split("x"))
            if None in (x, y, width, height):
                return None
            return x, y, width, height
        except FileNotFoundError:
            return None

    def set_phase(self, phase: str):
        if phase == "takeoff":
            self.start_time = time.time()
        elif phase == "search":
            self.search_start = time.time()
        elif phase != "search" and phase != "takeoff" and self.search_start > 0 and self.search_duration == 0:
            self.search_duration = time.time() - self.search_start

    def log_position(self, uav_x=0, uav_y=0, uav_alt=0, uav_heading=0, ugv_x=0, ugv_y=0):
        pass

    def start_logging(self):
        sys.stdout = self._tee_stdout
        sys.stderr = self._tee_stderr

    def stop_logging(self):
        sys.stdout = self._original_stdout
        sys.stderr = self._original_stderr

    def start_recording(self, tracker=None, camera=None):
        self._tracker_ref = tracker
        self._camera_ref = camera
        self._recording = True

        window_id = self._find_gazebo_window_ffmpeg()
        geo = None
        if window_id:
            geo = self._get_window_geometry(window_id)

        if geo:
            x, y, width, height = geo
            width += width % 2
            height += height % 2
        else:
            x, y, width, height = 0, 0, 1920, 1080

        self.gazebo_output_file = str(self.run_dir / "gazebo.mp4")
        display = os.environ.get("DISPLAY", ":0")
        if not display.endswith(".0"):
            display += ".0"

        cmd = [
            "ffmpeg", "-y",
            "-f", "x11grab"
        ]

        if window_id:
            cmd.extend([
                "-window_id", str(window_id),
                "-r", str(self.fps),
                "-i", display,
            ])
        else:
            cmd.extend([
                "-r", str(self.fps),
                "-s", f"{width}x{height}",
                "-i", f"{display}+{x},{y}",
            ])

        cmd.extend([
            "-c:v", "libx264",
            "-preset", "ultrafast",
            "-crf", "18",
            "-pix_fmt", "yuv420p",
            "-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2",
            "-movflags", "+faststart",
            self.gazebo_output_file
        ])
        try:
            self._gazebo_ffmpeg_proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            print(f"[REC] Started recording Gazebo window via ffmpeg to gazebo.mp4 (ID: {window_id})")
        except Exception as e:
            print(f"[REC] Failed to start ffmpeg: {e}")
            self._gazebo_ffmpeg_proc = None

        self._record_thread = threading.Thread(target=self._record_loop, daemon=True)
        self._record_thread.start()

    def _record_loop(self):
        interval = 1.0 / self.fps
        while self._recording:
            t0 = time.time()

            if self._tracker_ref is not None:
                try:
                    frame = self._tracker_ref.draw()
                    if frame is not None:
                        self._write_tracker_frame(frame)
                        self._last_tracker_frame = frame.copy()
                except Exception:
                    pass

            if self._camera_ref is not None:
                try:
                    frame = self._camera_ref.frames.get("downward")
                    if frame is not None:
                        self._write_camera_frame(frame)
                        self._last_camera_frame = frame.copy()
                except Exception:
                    pass

                try:
                    ugv_frame = self._camera_ref.frames.get("ugv_forward")
                    if ugv_frame is not None:
                        self._write_ugv_camera_frame(ugv_frame)
                        self._last_ugv_camera_frame = ugv_frame.copy()
                except Exception:
                    pass

            elapsed = time.time() - t0
            sleep_time = max(0, interval - elapsed)
            time.sleep(sleep_time)

    def _write_tracker_frame(self, frame):
        h, w = frame.shape[:2]
        if self._tracker_writer is None:
            self._tracker_size = (w, h)
            fourcc = cv2.VideoWriter_fourcc(*'XVID')
            path = str(self.run_dir / "flight_path.avi")
            self._tracker_writer = cv2.VideoWriter(path, fourcc, self.fps, (w, h))
        self._tracker_writer.write(frame)
        self._tracker_frames += 1

    def _write_camera_frame(self, frame):
        h, w = frame.shape[:2]
        if self._camera_writer is None:
            self._camera_size = (w, h)
            fourcc = cv2.VideoWriter_fourcc(*'XVID')
            path = str(self.run_dir / "camera.avi")
            self._camera_writer = cv2.VideoWriter(path, fourcc, self.fps, (w, h))
        self._camera_writer.write(frame)
        self._camera_frames += 1

    def _write_ugv_camera_frame(self, frame):
        h, w = frame.shape[:2]
        if self._ugv_camera_writer is None:
            self._ugv_camera_size = (w, h)
            fourcc = cv2.VideoWriter_fourcc(*'XVID')
            path = str(self.run_dir / "ugv_camera.avi")
            self._ugv_camera_writer = cv2.VideoWriter(path, fourcc, self.fps, (w, h))
        self._ugv_camera_writer.write(frame)
        self._ugv_camera_frames += 1

    def snapshot_camera(self, label="snapshot"):
        if self._camera_ref is None:
            return
        frame = self._camera_ref.frames.get("downward")
        if frame is None:
            return
        idx = len(self._camera_snapshots)
        filename = f"camera_{idx:03d}_{label}.png"
        path = self.run_dir / filename
        cv2.imwrite(str(path), frame)
        self._camera_snapshots.append(filename)
        print(f"[REC] Snapshot: {filename}")
        threading.Thread(target=self._upload_file, args=(path, filename), daemon=True).start()

    def snapshot_tracker(self, label="tracker"):
        if self._last_tracker_frame is not None:
            filename = f"tracker_{label}.png"
            path = self.run_dir / filename
            cv2.imwrite(str(path), self._last_tracker_frame)
            print(f"[REC] Tracker snapshot: {filename}")
            threading.Thread(target=self._upload_file, args=(path, filename), daemon=True).start()

    def save_summary(self, search_mode="", altitude=0, markers=None, landed=False, ugv_dispatched=False, ugv_target=None, extra=None):
        duration = time.time() - self.start_time
        summary = {
            "search_mode": search_mode,
            "altitude": round(altitude, 2),
            "duration_seconds": round(duration, 1),
            "search_duration_seconds": round(self.search_duration, 1),
            "landed": landed,
            "ugv_dispatched": ugv_dispatched,
            "markers_found": 0,
            "markers": {},
            "timestamp": datetime.now().isoformat(),
        }
        if ugv_target:
            summary["ugv_target"] = ugv_target
        if markers:
            summary["markers_found"] = len(markers)
            for mid, info in markers.items():
                pos = info.get("uav_position", {})
                summary["markers"][int(mid)] = {
                    "x": round(pos.get("x", 0), 2),
                    "y": round(pos.get("y", 0), 2),
                    "distance": round(info.get("distance", 0), 2),
                }
        if extra:
            summary.update(extra)

        path = self.run_dir / "summary.yaml"
        try:
            with open(path, "w") as f:
                yaml.dump(summary, f, default_flow_style=False, sort_keys=False)
            print(f"[REC] Summary saved: {path}")
            threading.Thread(target=self._upload_file, args=(path, "summary.yaml"), daemon=True).start()
        except Exception as e:
            print(f"[REC] Failed to save summary: {e}")

    def stop(self):
        self._recording = False
        if self._record_thread:
            self._record_thread.join(timeout=3.0)

        if hasattr(self, '_gazebo_ffmpeg_proc') and self._gazebo_ffmpeg_proc:
            try:
                self._gazebo_ffmpeg_proc.send_signal(signal.SIGINT)
                self._gazebo_ffmpeg_proc.wait(timeout=5)
            except Exception:
                pass
            self._upload_file(self.gazebo_output_file, "gazebo.mp4")

        if self._last_tracker_frame is not None:
            filename = "flight_path.png"
            path = self.run_dir / filename
            cv2.imwrite(str(path), self._last_tracker_frame)
            print(f"[REC] Flight path saved: {path}")
            self._upload_file(path, filename)

        if self._last_camera_frame is not None:
            filename = "camera_final.png"
            path = self.run_dir / filename
            cv2.imwrite(str(path), self._last_camera_frame)
            self._upload_file(path, filename)

        if self._last_ugv_camera_frame is not None:
            filename = "ugv_camera_final.png"
            path = self.run_dir / filename
            cv2.imwrite(str(path), self._last_ugv_camera_frame)
            self._upload_file(path, filename)

        if self._tracker_writer:
            self._tracker_writer.release()
            self._upload_file(self.run_dir / "flight_path.avi", "flight_path.avi")
        if self._camera_writer:
            self._camera_writer.release()
            self._upload_file(self.run_dir / "camera.avi", "camera.avi")
        if self._ugv_camera_writer:
            self._ugv_camera_writer.release()
            self._upload_file(self.run_dir / "ugv_camera.avi", "ugv_camera.avi")

        self.stop_logging()
        self._log_file.close()
        self._upload_file(self._log_path, "log.txt")

        duration = time.time() - self.start_time
        mins = int(duration // 60)
        secs = int(duration % 60)

        if self.sim_id:
            try:
                requests.put(
                    f"{API_URL}/api/simulations/{self.sim_id}/time",
                    json={"search_time": self.search_duration, "total_time": duration},
                    timeout=5
                )
            except Exception as e:
                print(f"[REC] PUT time failed: {e}")

        self._original_stdout.write(
            f"\n[REC] ═══════════════════════════════════════\n"
            f"[REC]   Cloud Upload Complete (ID: {self.sim_id})\n"
            f"[REC]   Duration: {mins}m {secs}s\n"
            f"[REC]   Tracker:    {self._tracker_frames} frames\n"
            f"[REC]   UAV Camera: {self._camera_frames} frames\n"
            f"[REC]   UGV Camera: {self._ugv_camera_frames} frames\n"
            f"[REC]   Gazebo:     via ffmpeg (.mp4)\n"
            f"[REC] ═══════════════════════════════════════\n"
        )

class _TeeWriter:
    def __init__(self, stream, log_file):
        self._stream = stream
        self._log = log_file

    def write(self, data):
        self._stream.write(data)
        try:
            self._log.write(data)
            self._log.flush()
        except (ValueError, IOError):
            pass

    def flush(self):
        self._stream.flush()
        try:
            self._log.flush()
        except (ValueError, IOError):
            pass

    def fileno(self):
        return self._stream.fileno()

    def isatty(self):
        return self._stream.isatty()