import os
import json
import logging
import datetime
import base64
import re
import requests
import paho.mqtt.client as mqtt
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from models.models import Client, ClientLog, LogLevel, ProcessStatus, ScreenHealthStatus
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)s] %(message)s')

# Load .env only when not already configured by Docker (API_BASE_URL not set by compose means we're outside a container)
_api_already_set = bool(os.environ.get("API_BASE_URL"))
if not _api_already_set and os.getenv("ENV", "development") == "development":
    try:
        from dotenv import load_dotenv
        load_dotenv(".env")
    except Exception:
        pass

# ENV-dependent configuration
ENV = os.getenv("ENV", "development")
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO" if ENV == "production" else "DEBUG")
DB_URL = os.environ.get("DB_CONN", "mysql+pymysql://user:password@db/infoscreen")

# DB-Konfiguration
engine = create_engine(DB_URL)
Session = sessionmaker(bind=engine)

# API configuration
API_BASE_URL = os.getenv("API_BASE_URL", "http://server:8000")


def normalize_process_status(value):
    if value is None:
        return None
    if isinstance(value, ProcessStatus):
        return value

    normalized = str(value).strip().lower()
    if not normalized:
        return None

    try:
        return ProcessStatus(normalized)
    except ValueError:
        return None


def normalize_event_id(value):
    if value is None or isinstance(value, bool):
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value)

    normalized = str(value).strip()
    if not normalized:
        return None
    if normalized.isdigit():
        return int(normalized)

    match = re.search(r"(\d+)$", normalized)
    if match:
        return int(match.group(1))

    return None


def parse_timestamp(value):
    if not value:
        return None
    if isinstance(value, (int, float)):
        try:
            ts_value = float(value)
            if ts_value > 1e12:
                ts_value = ts_value / 1000.0
            return datetime.datetime.fromtimestamp(ts_value, datetime.UTC)
        except (TypeError, ValueError, OverflowError):
            return None
    try:
        value_str = str(value).strip()
        if value_str.isdigit():
            ts_value = float(value_str)
            if ts_value > 1e12:
                ts_value = ts_value / 1000.0
            return datetime.datetime.fromtimestamp(ts_value, datetime.UTC)

        parsed = datetime.datetime.fromisoformat(value_str.replace('Z', '+00:00'))
        if parsed.tzinfo is None:
            return parsed.replace(tzinfo=datetime.UTC)
        return parsed.astimezone(datetime.UTC)
    except ValueError:
        return None


def infer_screen_health_status(payload_data):
    explicit = payload_data.get('screen_health_status')
    if explicit:
        try:
            return ScreenHealthStatus[str(explicit).strip().upper()]
        except KeyError:
            pass

    metrics = payload_data.get('health_metrics') or {}
    if metrics.get('screen_on') is False:
        return ScreenHealthStatus.BLACK

    last_frame_update = parse_timestamp(metrics.get('last_frame_update'))
    if last_frame_update:
        age_seconds = (datetime.datetime.now(datetime.UTC) - last_frame_update).total_seconds()
        if age_seconds > 30:
            return ScreenHealthStatus.FROZEN
        return ScreenHealthStatus.OK

    return None


def apply_monitoring_update(client_obj, *, event_id=None, process_name=None, process_pid=None,
                            process_status=None, last_seen=None, screen_health_status=None,
                            last_screenshot_analyzed=None):
    if last_seen:
        client_obj.last_alive = last_seen

    normalized_event_id = normalize_event_id(event_id)
    if normalized_event_id is not None:
        client_obj.current_event_id = normalized_event_id

    if process_name is not None:
        client_obj.current_process = process_name

    if process_pid is not None:
        client_obj.process_pid = process_pid

    normalized_status = normalize_process_status(process_status)
    if normalized_status is not None:
        client_obj.process_status = normalized_status

    if screen_health_status is not None:
        client_obj.screen_health_status = screen_health_status

    if last_screenshot_analyzed is not None:
        existing = client_obj.last_screenshot_analyzed
        if existing is not None and existing.tzinfo is None:
            existing = existing.replace(tzinfo=datetime.UTC)

        candidate = last_screenshot_analyzed
        if candidate.tzinfo is None:
            candidate = candidate.replace(tzinfo=datetime.UTC)

        if existing is None or candidate >= existing:
            client_obj.last_screenshot_analyzed = candidate


def _extract_image_and_timestamp(data):
    image_value = None
    timestamp_value = None
    screenshot_type = None

    if not isinstance(data, dict):
        return None, None, None

    screenshot_obj = data.get("screenshot") if isinstance(data.get("screenshot"), dict) else None
    metadata_obj = data.get("metadata") if isinstance(data.get("metadata"), dict) else None
    screenshot_meta_obj = screenshot_obj.get("metadata") if screenshot_obj and isinstance(screenshot_obj.get("metadata"), dict) else None

    for container in (data, screenshot_obj, metadata_obj, screenshot_meta_obj):
        if not isinstance(container, dict):
            continue
        raw_type = container.get("screenshot_type") or container.get("screenshotType")
        if raw_type is not None:
            normalized_type = str(raw_type).strip().lower()
            if normalized_type in ("periodic", "event_start", "event_stop"):
                screenshot_type = normalized_type
                break

    for key in ("image", "data"):
        if isinstance(data.get(key), str) and data.get(key):
            image_value = data.get(key)
            break
    if image_value is None and screenshot_obj is not None:
        for key in ("image", "data"):
            if isinstance(screenshot_obj.get(key), str) and screenshot_obj.get(key):
                image_value = screenshot_obj.get(key)
                break

    for container in (data, screenshot_obj, metadata_obj, screenshot_meta_obj):
        if not isinstance(container, dict):
            continue
        for key in ("timestamp", "captured_at", "capture_time", "created_at"):
            value = container.get(key)
            if value is not None:
                timestamp_value = value
                return image_value, timestamp_value, screenshot_type

    return image_value, timestamp_value, screenshot_type


def handle_screenshot(uuid, payload):
    """
    Handle screenshot data received via MQTT and forward to API.
    Payload can be either raw binary image data or JSON with base64-encoded image.
    """
    try:
        # Try to parse as JSON first
        try:
            data = json.loads(payload.decode())
            image_b64, timestamp_value, screenshot_type = _extract_image_and_timestamp(data)
            if image_b64:
                # Payload is JSON with base64 image
                api_payload = {"image": image_b64}
                if timestamp_value is not None:
                    api_payload["timestamp"] = timestamp_value
                if screenshot_type:
                    api_payload["screenshot_type"] = screenshot_type
                headers = {"Content-Type": "application/json"}
                logging.debug(f"Forwarding base64 screenshot from {uuid} to API")
            else:
                logging.warning(f"Screenshot JSON from {uuid} missing image/data field")
                return
        except (json.JSONDecodeError, UnicodeDecodeError):
            # Payload is raw binary image data - encode to base64 for API
            image_b64 = base64.b64encode(payload).decode('utf-8')
            api_payload = {"image": image_b64}
            headers = {"Content-Type": "application/json"}
            logging.debug(f"Forwarding binary screenshot from {uuid} to API (encoded as base64)")

        # Forward to API endpoint
        api_url = f"{API_BASE_URL}/api/clients/{uuid}/screenshot"
        response = requests.post(api_url, json=api_payload, headers=headers, timeout=10)

        if response.status_code == 200:
            logging.info(f"Screenshot von {uuid} erfolgreich an API weitergeleitet")
        else:
            logging.error(f"API returned status {response.status_code} for screenshot from {uuid}: {response.text}")

    except requests.exceptions.RequestException as e:
        logging.error(f"Failed to forward screenshot from {uuid} to API: {e}")
    except Exception as e:
        logging.error(f"Error handling screenshot from {uuid}: {e}")


def on_connect(client, userdata, flags, reasonCode, properties):
    """Callback for when client connects or reconnects (API v2)."""
    try:
        # Subscribe on every (re)connect so we don't miss heartbeats after broker restarts
        client.subscribe("infoscreen/discovery")
        client.subscribe("infoscreen/+/heartbeat")
        client.subscribe("infoscreen/+/screenshot")
        client.subscribe("infoscreen/+/dashboard")
        
        # Subscribe to monitoring topics
        client.subscribe("infoscreen/+/logs/error")
        client.subscribe("infoscreen/+/logs/warn")
        client.subscribe("infoscreen/+/logs/info")
        client.subscribe("infoscreen/+/health")
        
        logging.info(f"MQTT connected (reasonCode: {reasonCode}); (re)subscribed to discovery, heartbeats, screenshots, dashboards, logs, and health")
    except Exception as e:
        logging.error(f"Subscribe failed on connect: {e}")


def on_message(client, userdata, msg):
    topic = msg.topic
    logging.debug(f"Empfangene Nachricht auf Topic: {topic}")

    try:
        # Dashboard-Handling (nested screenshot payload)
        if topic.startswith("infoscreen/") and topic.endswith("/dashboard"):
            uuid = topic.split("/")[1]
            try:
                payload_text = msg.payload.decode()
                data = json.loads(payload_text)
                image_b64, ts_value, screenshot_type = _extract_image_and_timestamp(data)
                if image_b64:
                    logging.debug(f"Dashboard enthält Screenshot für {uuid}; Weiterleitung an API")
                    dashboard_payload = {"image": image_b64}
                    if ts_value is not None:
                        dashboard_payload["timestamp"] = ts_value
                    if screenshot_type:
                        dashboard_payload["screenshot_type"] = screenshot_type
                    api_payload = json.dumps(dashboard_payload).encode("utf-8")
                    handle_screenshot(uuid, api_payload)
                # Update last_alive if status present
                if data.get("status") == "alive":
                    session = Session()
                    client_obj = session.query(Client).filter_by(uuid=uuid).first()
                    if client_obj:
                        process_health = data.get('process_health') or {}
                        apply_monitoring_update(
                            client_obj,
                            last_seen=datetime.datetime.now(datetime.UTC),
                            event_id=process_health.get('event_id'),
                            process_name=process_health.get('current_process') or process_health.get('process'),
                            process_pid=process_health.get('process_pid') or process_health.get('pid'),
                            process_status=process_health.get('process_status') or process_health.get('status'),
                        )
                        session.commit()
                    session.close()
            except Exception as e:
                logging.error(f"Fehler beim Verarbeiten des Dashboard-Payloads von {uuid}: {e}")
            return

        # Screenshot-Handling
        if topic.startswith("infoscreen/") and topic.endswith("/screenshot"):
            uuid = topic.split("/")[1]
            handle_screenshot(uuid, msg.payload)
            return

        # Heartbeat-Handling
        if topic.startswith("infoscreen/") and topic.endswith("/heartbeat"):
            uuid = topic.split("/")[1]
            try:
                # Parse payload to get optional health data
                payload_data = json.loads(msg.payload.decode())
            except (json.JSONDecodeError, UnicodeDecodeError):
                payload_data = {}
            
            session = Session()
            client_obj = session.query(Client).filter_by(uuid=uuid).first()
            if client_obj:
                apply_monitoring_update(
                    client_obj,
                    last_seen=datetime.datetime.now(datetime.UTC),
                    event_id=payload_data.get('current_event_id'),
                    process_name=payload_data.get('current_process'),
                    process_pid=payload_data.get('process_pid'),
                    process_status=payload_data.get('process_status'),
                )
                session.commit()
                logging.info(f"Heartbeat von {uuid} empfangen, last_alive (UTC) aktualisiert.")
            session.close()
            return
        
        # Log-Handling (ERROR, WARN, INFO)
        if topic.startswith("infoscreen/") and "/logs/" in topic:
            parts = topic.split("/")
            if len(parts) >= 4:
                uuid = parts[1]
                level_str = parts[3].upper()  # 'error', 'warn', 'info' -> 'ERROR', 'WARN', 'INFO'
                
                try:
                    payload_data = json.loads(msg.payload.decode())
                    message = payload_data.get('message', '')
                    timestamp_str = payload_data.get('timestamp')
                    context = payload_data.get('context', {})
                    
                    # Parse timestamp or use current time
                    if timestamp_str:
                        try:
                            log_timestamp = datetime.datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
                            if log_timestamp.tzinfo is None:
                                log_timestamp = log_timestamp.replace(tzinfo=datetime.UTC)
                        except ValueError:
                            log_timestamp = datetime.datetime.now(datetime.UTC)
                    else:
                        log_timestamp = datetime.datetime.now(datetime.UTC)
                    
                    # Store in database
                    session = Session()
                    try:
                        log_level = LogLevel[level_str]
                        log_entry = ClientLog(
                            client_uuid=uuid,
                            timestamp=log_timestamp,
                            level=log_level,
                            message=message,
                            context=json.dumps(context) if context else None
                        )
                        session.add(log_entry)
                        session.commit()
                        logging.info(f"[{level_str}] {uuid}: {message}")
                    except Exception as e:
                        logging.error(f"Error saving log from {uuid}: {e}")
                        session.rollback()
                    finally:
                        session.close()
                    
                except (json.JSONDecodeError, UnicodeDecodeError) as e:
                    logging.error(f"Could not parse log payload from {uuid}: {e}")
            return
        
        # Health-Handling
        if topic.startswith("infoscreen/") and topic.endswith("/health"):
            uuid = topic.split("/")[1]
            try:
                payload_data = json.loads(msg.payload.decode())
                
                session = Session()
                client_obj = session.query(Client).filter_by(uuid=uuid).first()
                if client_obj:
                    # Update expected state
                    expected = payload_data.get('expected_state', {})

                    # Update actual state
                    actual = payload_data.get('actual_state', {})
                    screen_health_status = infer_screen_health_status(payload_data)
                    apply_monitoring_update(
                        client_obj,
                        last_seen=datetime.datetime.now(datetime.UTC),
                        event_id=expected.get('event_id'),
                        process_name=actual.get('process'),
                        process_pid=actual.get('pid'),
                        process_status=actual.get('status'),
                        screen_health_status=screen_health_status,
                        last_screenshot_analyzed=parse_timestamp((payload_data.get('health_metrics') or {}).get('last_frame_update')),
                    )
                    session.commit()
                    logging.debug(f"Health update from {uuid}: {actual.get('process')} ({actual.get('status')})")
                session.close()
                
            except (json.JSONDecodeError, UnicodeDecodeError) as e:
                logging.error(f"Could not parse health payload from {uuid}: {e}")
            except Exception as e:
                logging.error(f"Error processing health from {uuid}: {e}")
            return

        # Discovery-Handling
        if topic == "infoscreen/discovery":
            payload = json.loads(msg.payload.decode())
            logging.info(f"Discovery empfangen: {payload}")

            if "uuid" in payload:
                uuid = payload["uuid"]
                session = Session()
                existing = session.query(Client).filter_by(uuid=uuid).first()

                if not existing:
                    new_client = Client(
                        uuid=uuid,
                        hardware_token=payload.get("hardware_token"),
                        ip=payload.get("ip"),
                        type=payload.get("type"),
                        hostname=payload.get("hostname"),
                        os_version=payload.get("os_version"),
                        software_version=payload.get("software_version"),
                        macs=",".join(payload.get("macs", [])),
                        model=payload.get("model"),
                        registration_time=datetime.datetime.now(datetime.UTC),
                    )
                    session.add(new_client)
                    session.commit()
                    logging.info(f"Neuer Client registriert: {uuid}")
                else:
                    logging.info(f"Client bereits bekannt: {uuid}")

                session.close()

                # Discovery-ACK senden
                ack_topic = f"infoscreen/{uuid}/discovery_ack"
                client.publish(ack_topic, json.dumps({"status": "ok"}))
                logging.info(f"Discovery-ACK gesendet an {ack_topic}")
            else:
                logging.warning("Discovery ohne UUID empfangen, ignoriert.")

    except Exception as e:
        logging.error(f"Fehler bei Verarbeitung: {e}")


def main():
    mqtt_client = mqtt.Client(protocol=mqtt.MQTTv311, callback_api_version=mqtt.CallbackAPIVersion.VERSION2)
    mqtt_client.on_message = on_message
    mqtt_client.on_connect = on_connect
    # Set an exponential reconnect delay to survive broker restarts
    mqtt_client.reconnect_delay_set(min_delay=1, max_delay=60)
    mqtt_client.connect("mqtt", 1883)

    logging.info("Listener gestartet; warte auf MQTT-Verbindung und Nachrichten")
    mqtt_client.loop_forever()


if __name__ == "__main__":
    main()
