import os
import sys
from typing import Optional

import av
import torch
import streamlit as st
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, RTCConfiguration

try:
    from ultralytics import YOLO
except Exception as e:
    st.error(f"Ultralytics import failed: {e}")
    raise


st.set_page_config(page_title="YOLO11 Webcam", layout="wide")
st.title("YOLO11 Webcam Demo")
st.caption("Live object detection with YOLO11 using your webcam. Use the sidebar to configure the model and thresholds.")


def _resolve_device(device: str) -> str:
    """Map UI device option to a valid torch/Ultralytics device string."""
    d = (device or "").lower().strip()
    if d in {"cpu", "cuda", "mps", "xpu", "dml"}:
        return d
    if d == "auto":
        if torch.cuda.is_available():
            return "cuda"
        # Apple MPS (macOS). Kept for completeness even if not typical on Windows.
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
        # Windows DirectML (requires torch-directml)
        try:
            import torch_directml  # noqa: F401
            return "dml"
        except Exception:
            pass
        return "cpu"
    # Fallback to CPU for any unknown value
    return "cpu"


@st.cache_resource(show_spinner=True)
def load_model(model_source: str, device: str = "auto"):
    """
    Load an Ultralytics model from a local path, built-in alias, or Hugging Face Hub path.
    Examples for model_source:
      - "yolo11n.pt" (auto-downloads if available)
      - "C:/path/to/your_model.pt"
      - "hf://Ultralytics/YOLO11/yolo11n.pt" (HF Hub file)
    """
    # Resolve device (maps 'auto' to an actual device)
    resolved_device = _resolve_device(device)

    # Gentle hint for common mistake: loading a multi-model container instead of a single YOLO weight
    try:
        base_name = os.path.basename(model_source)
        if base_name.lower().startswith("models_multi"):
            st.info("Seçilen .pt birden fazla ağırlık içeren bir paket olabilir. Lütfen tek bir YOLO ağırlık dosyası seçin (ör. yolo11n.pt).")
    except Exception:
        pass

    # 1) Kaynağı yerel dosyaya indir/çöz ve YOLO nesnesini oluştur
    def _create_yolo(path_or_alias: str):
        if path_or_alias.startswith("hf://"):
            try:
                from huggingface_hub import hf_hub_download
            except Exception:
                st.error("huggingface_hub yüklü değil. Yükleyin veya yerel/yerleşik model kullanın.")
                raise
            path = path_or_alias.replace("hf://", "", 1)
            if "/" not in path:
                raise ValueError("hf:// için hf://<repo_id>/<filename> biçimini kullanın")
            repo_id, filename = path.split("/", 1)
            st.info(f"{repo_id} deposundan {filename} indiriliyor…")
            local_path = hf_hub_download(repo_id=repo_id, filename=filename, token=os.getenv("HF_TOKEN"))
            return YOLO(local_path)
        else:
            return YOLO(path_or_alias)

    # Önce modeli yüklemeyi deneyin (hata ==> dosya/format sorunu)
    try:
        model = _create_yolo(model_source)
    except Exception as e:
        st.error(
            "Model dosyası yüklenemedi. Bu dosya geçerli bir YOLO ağırlığı olmayabilir (örn. eğitim checkpoint'i veya paketlenmiş container).\n"
            f"Detay: {e}"
        )
        raise

    # 2) Sonra cihaza taşımayı deneyin (hata ==> CUDA/MPS taşınamıyor → CPU'ya düş)
    try:
        return model.to(resolved_device)
    except Exception as e:
        st.warning(f"{resolved_device} cihaza taşınamadı, CPU'ya düşülüyor. Detay: {e}")
        try:
            return model.to("cpu")
        except Exception:
            # CPU'ya taşınma da başarısızsa, ham modeli döndür (son çare)
            return model


class YOLOProcessor(VideoProcessorBase):
    def __init__(self, model: YOLO, conf: float, iou: float):
        self.model = model
        self.conf = conf
        self.iou = iou

    def recv(self, frame: av.VideoFrame) -> av.VideoFrame:
        img = frame.to_ndarray(format="bgr24")
        results = self.model.predict(img, conf=self.conf, iou=self.iou, verbose=False)
        plotted = results[0].plot()
        return av.VideoFrame.from_ndarray(plotted, format="bgr24")


with st.sidebar:
    st.header("Ayarlar")
    default_model = "yolo11n.pt"  # Ultralytics dağıtımından otomatik indirilmeye çalışılır
    model_path = st.text_input("Model yolu veya alias", value=default_model, help="Yerel .pt yolu, yerleşik alias (örn. yolo11n.pt) veya hf://Ultralytics/YOLO11/yolo11n.pt")
    device = st.selectbox("Cihaz", options=["auto", "cpu", "cuda", "dml"], index=0)
    conf = st.slider("Confidence", min_value=0.1, max_value=0.9, value=0.25, step=0.05)
    iou = st.slider("IoU", min_value=0.1, max_value=0.9, value=0.45, step=0.05)
    load_btn = st.button("Modeli Yükle")

    # Ortam durumu (Torch/CUDA)
    try:
        tv = torch.__version__
        cv = getattr(torch.version, "cuda", None)
        ca = torch.cuda.is_available()
        # Detect DirectML
        has_dml = False
        try:
            import torch_directml  # noqa: F401
            has_dml = True
        except Exception:
            has_dml = False
        dml_txt = " • DirectML: Hazır" if has_dml else ""
        st.caption(f"Torch {tv} • CUDA: {cv or 'yok'} • GPU etkin (CUDA): {'Evet' if ca else 'Hayır'}{dml_txt}")
    except Exception:
        pass


st.session_state.setdefault("model", None)


if load_btn:
    try:
        st.session_state.model = load_model(model_path, device=device)
        names = getattr(st.session_state.model, "names", None)
        if isinstance(names, dict):
            st.success(f"Model yüklendi. Sınıflar: {len(names)}")
        else:
            st.success("Model yüklendi.")
        try:
            active_device = next((p.device.type for p in st.session_state.model.parameters()), "unknown")
            st.info(f"Aktif cihaz: {active_device}")
        except Exception:
            pass
    except Exception as e:
        st.exception(e)


rtc_config = RTCConfiguration({
    "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}],
})


col1, col2 = st.columns([2, 1])
with col1:
    st.subheader("Webcam")
    if st.session_state.get("model") is None:
        st.info("Önce soldan bir model yükleyin.")
    else:
        # Capture current values to avoid session_state race during worker creation
        _model = st.session_state.get("model")
        _conf = float(conf)
        _iou = float(iou)

        def video_processor_factory(model=_model, conf_val=_conf, iou_val=_iou):
            return YOLOProcessor(model, conf=conf_val, iou=iou_val)

        webrtc_streamer(
            key="yolo11-webcam",
            video_processor_factory=video_processor_factory,
            rtc_configuration=rtc_config,
            media_stream_constraints={"video": True, "audio": False},
        )

with col2:
    st.subheader("İpuçları")
    st.markdown(
        """
        - Model alanına şunlardan birini girebilirsiniz:
            - Yerleşik: `yolo11n.pt` (veya sizde olan başka bir .pt)
            - Yerel dosya: `C:/yolov11/weights/custom.pt`
            - HF Hub: `hf://Ultralytics/YOLO11/yolo11n.pt`
        - Hugging Face özel/korumalı dosyalar için `HF_TOKEN` ortam değişkenini ayarlayabilirsiniz.
        - GPU (CUDA) yoksa cihazı `cpu` bırakabilirsiniz.
        - Stream durmuyorsa tarayıcı izinlerini kontrol edin ve sayfayı yenileyin.
        """
    )