#!/usr/bin/env python3
"""
I2C Doctor — Automated Generic Sensor Qualification Suite.

Version: V1.0.0

Tests any I2C sensor thoroughly regarding electrical health, maximum speed,
long cable attenuation tolerance, and read reliability.

Usage:
    python3 auto_test_sensor.py --port /dev/tty.usbserial-* [--addr 0x48] [--speed 1000]
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime
from bridge import I2CBridge

class SensorTestSuite:
    def __init__(self, port: str, baud: int = 115200, target_addr: int = None, 
                 target_reg: int = None, test_speed: int = 100):
        self.port = port
        self.baud = baud
        self.target_addr = target_addr
        self.target_reg = target_reg
        self.test_speed = test_speed
        self.report = {
            "timestamp": datetime.now().isoformat(),
            "target_addr": None,
            "passed": False,
            "test_speed_khz": self.test_speed,
            "summary": {},
            "diagnostics": {},
            "speed_tests": {},
            "stress_tests": {},
            "fault_injection": {},
            "register_dump": {}
        }

    def log(self, section: str, msg: str, level: str = "INFO"):
        icon = {"INFO": "ℹ️", "SUCCESS": "✅", "WARN": "⚠️", "FAIL": "❌", "TEST": "🧪"}.get(level, " ")
        print(f"[{datetime.now().strftime('%H:%M:%S')}] {icon} [{section}] {msg}")

    def run_all(self) -> bool:
        print("\n" + "="*70)
        print("    I2C DOCTOR — AUTOMATED SENSOR QUALIFICATION SUITE    ")
        print("="*70 + "\n")
        try:
            with I2CBridge(self.port, self.baud) as b:
                # 0. Reset everything
                self.log("INIT", "Resetting hardware & clearing bus state...", "INFO")
                self._reset_hardware(b)
                b.bus_recovery()
                time.sleep(0.05)
                # 1. Bus diagnostics
                if not self._test_bus_diagnostics(b):
                    self.log("BUS", "Electrical bus diagnostics failed! Aborting.", "FAIL")
                    return False
                # 2. Bus scan & target identification
                if not self._test_scan_and_probe(b):
                    self.log("SCAN", "No reachable I2C device found! Aborting.", "FAIL")
                    return False
                # 3. Speed tolerance test (10 Samples per Step)
                self._test_speed_capabilities(b, samples=10)
                # 4. Stress test
                self._test_stress_limits(b)
                # 5. Cable Attenuation Test
                self._test_fault_tolerance(b)
                # 6. Optional register dump
                self._dump_registers(b)
                # Evaluate final overall status
                stress_rate = self._get_stress_pass_rate()
                self.report["passed"] = (stress_rate >= 99.0)
                self._print_summary()
                return self.report["passed"]
        except Exception as e:
            self.log("CRITICAL", f"Unexpected error during testing: {e}", "FAIL")
            return False

    def _reset_hardware(self, b: I2CBridge):
        b.long_cable(0)
        b.noise(0)
        b.sda_hold(0)
        b.clock_stretch(0)
        b.speed(self.test_speed)
        time.sleep(0.02)

    def _test_bus_diagnostics(self, b: I2CBridge) -> bool:
        self.log("DIAG", "Starting electrical bus analysis...", "TEST")
        diag = b.diagnose()
        self.report["diagnostics"] = diag
        vdd = diag.get("vdd", 0)
        score = diag.get("score", 0)
        shorts = diag.get("shorts", False)
        self.log("DIAG", f"VDD: {vdd} V | Bus Score: {score}% | Shorts: {shorts}")
        has_short = False
        if isinstance(shorts, bool):
            has_short = shorts
        elif isinstance(shorts, str):
            has_short = shorts.upper() in ("TRUE", "DETECTED", "YES")
        if has_short:
            self.log("DIAG", "Short circuit detected on SDA/SCL!", "FAIL")
            return False
        if score < 50:
            self.log("DIAG", f"Insufficient bus quality ({score}%). Check pull-up resistors.", "WARN")
        return True

    def _test_scan_and_probe(self, b: I2CBridge) -> bool:
        self.log("SCAN", "Scanning I2C bus...", "TEST")
        found = b.scan()
        # plausibility check: when > 30 devices we see phantom ACKS / SDA stuck LOW
        if len(found) > 30:
            self.log("SCAN", f"Ghost bus detected! Too many ACKs ({len(found)} devices). SDA might be held LOW.", "FAIL")
            return False
        if not found:
            self.log("SCAN", "Attempting I2C bus recovery...", "WARN")
            b.bus_recovery()
            found = b.scan()
        if not found or len(found) > 30:
            return False
        hex_addrs = [f"0x{a:02X}" for a in found]
        self.log("SCAN", f"Found I2C addresses: {', '.join(hex_addrs)}", "SUCCESS")
        if self.target_addr is not None:
            if self.target_addr in found:
                addr = self.target_addr
            else:
                self.log("SCAN", f"Target address 0x{self.target_addr:02X} not found on the bus!", "FAIL")
                return False
        else:
            # ignore invalid I2C addresses (0x00-0x07, 0x78-0x7F)
            valid_targets = [a for a in found if 0x08 <= a <= 0x77]
            if not valid_targets:
                self.log("SCAN", "No valid standard sensor address found (range 0x08..0x77)!", "FAIL")
                return False
            addr = valid_targets[0]
        self.target_addr = addr
        self.report["target_addr"] = f"0x{addr:02X}"
        self.log("SCAN", f"Selected target address for qualification: 0x{addr:02X}", "INFO")
        if b.probe(addr):
            self.log("PROBE", f"PROBE on 0x{addr:02X} successful (ACK)", "SUCCESS")
            return True
        else:
            self.log("PROBE", f"PROBE on 0x{addr:02X} failed (NACK)", "FAIL")
            return False

    def _sample_probe_read(self, b: I2CBridge, samples: int = 10) -> float:
        """Executes probe N times and returns the success rate in percent."""
        reg = None
        # check if a register read is required
        b.probe_read(self.target_addr, reg) # dummy probe_read in case first read is accepted but second is not
        if not b.probe_read(self.target_addr, reg):
            if b.probe_read(self.target_addr, 0):
                reg = 0x00
        acks = 0
        for _ in range(samples):
            if b.probe_read(self.target_addr, reg):
                acks += 1
        return (acks / samples) * 100.0

    def _test_speed_capabilities(self, b: I2CBridge, samples: int = 10):
        self.log("SPEED", f"Checking supported I2C clock speeds ({samples} samples per step)...", "TEST")
        speeds_to_test = [10, 50, 100, 200, 400, 1000]
        results = {}
        for spd in speeds_to_test:
            try:
                b.speed(spd)
                time.sleep(0.02)
                rate = self._sample_probe_read(b, samples=samples)
                # >= 99.0% to PASS
                status = "PASS" if rate >= 99.0 else ("DEGRADED" if rate > 0 else "FAIL")
                results[f"{spd}_kHz"] = {
                    "pass_rate": rate,
                    "status": status
                }               
                if status == "PASS":
                    self.log("SPEED", f"  {spd:4d} kHz: {rate:.1f}% ACKs [PASS]", "SUCCESS")
                elif status == "DEGRADED":
                    self.log("SPEED", f"  {spd:4d} kHz: {rate:.1f}% ACKs [DEGRADED]", "WARN")
                else:
                    self.log("SPEED", f"  {spd:4d} kHz: {rate:.1f}% ACKs [FAIL]", "WARN")
            except Exception as e:
                results[f"{spd}_kHz"] = {"pass_rate": 0.0, "status": f"ERROR ({e})"}
                self.log("SPEED", f"  {spd:4d} kHz: Error/Timeout", "WARN")
        # Restore requested test speed
        b.speed(self.test_speed)
        self.report["speed_tests"] = results

    def _test_stress_limits(self, b: I2CBridge):
        self.log("STRESS", f"Starting read stress test (1000 cycles @ {self.test_speed}kHz)...", "TEST")
        b.speed(self.test_speed)
        try:
            reg = self.target_reg
            if reg == None:
                # check if a register read is required
                if not b.probe_read(self.target_addr, reg):
                    if b.probe_read(self.target_addr, 0):
                        reg = 0x00
            res = b.stress_test(self.target_addr, count=1000, reg=reg)
            self.report["stress_tests"][f"{self.test_speed}kHz_1000_cycles"] = res
            pass_rate = res.get("pass_rate", 0.0)
            dt = res.get("dt", 0)
            self.log("STRESS", f"Result: {res['ack']}/{res['total']} ACKs ({pass_rate:.2f}%) in {dt} ms")
            if pass_rate >= 99.0: # >= 99.0% to PASS
                self.log("STRESS", "Stress test passed smoothly.", "SUCCESS")
            else:
                self.log("STRESS", f"Stress test showed packet losses! Pass rate: {pass_rate:.2f}%", "FAIL")
        except Exception as e:
            self.log("STRESS", f"Stress test aborted: {e}", "FAIL")

    def _test_fault_tolerance(self, b: I2CBridge):
        self.log("FAULT", "Starting physical bus attenuation tests...", "TEST")
        fault_results = {}
        self.log("FAULT", "  - Testing long cable / capacitance tolerance (LONG_CABLE 0 -> 5)...")
        cable_pass = 0
        for lvl in range(1, 6):
            b.long_cable(lvl)
            time.sleep(0.02)
            rate = self._sample_probe_read(b, samples=10)
            if rate >= 99.0: # >= 99.0% to PASS
                cable_pass = lvl
            else:
                break
        b.long_cable(0)
        time.sleep(0.02)
        fault_results["max_long_cable_level_passed"] = cable_pass
        self.log("FAULT", f"    Cable length tolerance: Level {cable_pass}/5", "INFO")
        self.report["fault_injection"] = fault_results

    def _dump_registers(self, b: I2CBridge):
        self.log("DUMP", "Reading first 16 registers (Reg 0x00 .. 0x0F)...", "TEST")
        reg_data = {}
        for reg in range(0x00, 0x10):
            try:
                val = b.read_reg(self.target_addr, reg, 1)
                reg_data[f"0x{reg:02X}"] = f"0x{val[0]:02X}"
            except Exception:
                reg_data[f"0x{reg:02X}"] = "N/A"
        self.report["register_dump"] = reg_data
        dump_str = " ".join([f"{k}:{v}" for k, v in list(reg_data.items())[:8]])
        self.log("DUMP", f"Registers [0x00..0x07]: {dump_str}", "INFO")

    def _get_stress_pass_rate(self) -> float:
        key = f"{self.test_speed}kHz_1000_cycles"
        stress_data = self.report["stress_tests"].get(key, {})
        return stress_data.get("pass_rate", 0.0)

    def _print_summary(self):
        stress_rate = self._get_stress_pass_rate()
        overall_passed = self.report["passed"]
        status_str = "QUALIFIED (PASS)" if overall_passed else "MARGINAL / UNSTABLE (FAIL)"
        print("\n" + "="*70)
        print(f"    QUALIFICATION RESULT FOR SENSOR 0x{self.target_addr:02X}    ")
        print("="*70)
        print(f" Status:                  {status_str}")
        print(f" Bus Electrical Score:    {self.report['diagnostics'].get('score', 'N/A')}%")
        print(f" Test Clock Speed:        {self.test_speed} kHz")
        print(f" Max Probe Speed:         {self._get_max_working_speed()}")
        print(f" Stress Test Pass Rate:   {stress_rate:.2f}% (1000 cycles @ {self.test_speed} kHz)")
        print(f" Cable Length Tolerance:  Level {self.report['fault_injection'].get('max_long_cable_level_passed', 0)} / 5")
        print("="*70 + "\n")

    def _get_max_working_speed(self) -> str:
        for spd in ["1000_kHz", "400_kHz", "200_kHz", "100_kHz", "50_kHz", "10_kHz"]:
            res = self.report["speed_tests"].get(spd, {})
            if res.get("status") == "PASS":
                return spd.replace("_", " ")
        return "None"

    def save_report_json(self, filename: str = None):
        output_dir = "reports"
        os.makedirs(output_dir, exist_ok=True)
        if not filename:
            addr_str = f"0x{self.target_addr:02X}" if self.target_addr else "unknown"
            filename = f"sensor_test_{addr_str}_{int(time.time())}.json"
        filepath = os.path.join(output_dir, os.path.basename(filename))
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(self.report, f, indent=2)
        self.log("FILE", f"Test report saved under: {filepath}", "SUCCESS")


def main():
    parser = argparse.ArgumentParser(description="Automated I2C Sensor Test Suite via I2C Doctor Bridge")
    parser.add_argument("--port", required=True, help="Serial port (e.g. /dev/tty.usbserial-110 or COM3)")
    parser.add_argument("--baud", type=int, default=115200, help="Baud rate (default: 115200)")
    parser.add_argument("--speed", type=int, default=100, help="Test I2C bus speed in kHz (default: 100)")
    parser.add_argument("--addr", help="Optional: Target I2C hex address (e.g. 0x48)")
    parser.add_argument("--reg", help="Optional: Specific test register in hex format (e.g. 0x00)")
    parser.add_argument("--json", help="Optional: Filename for saving the JSON report inside 'reports/' folder")
    args = parser.parse_args()
    target_addr = int(args.addr, 16) if args.addr else None
    target_reg = int(args.reg, 16) if args.reg else None
    suite = SensorTestSuite(
        port=args.port,
        baud=args.baud,
        target_addr=target_addr,
        target_reg=target_reg,
        test_speed=args.speed
    )
    success = suite.run_all()
    if args.json or success:
        suite.save_report_json(args.json)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()