import glob
import os
import logging as log
import uuid
import json
import sys
from datetime import datetime
import asyncio
import numpy as np
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from auracast import multicast_control, auracast_config
from aiortc import RTCPeerConnection, RTCSessionDescription, MediaStreamTrack
import av
import av.audio.layout
import sounddevice as sd  # type: ignore
from typing import Set
import traceback


PTIME = 40 # TODO: seems to have no effect at all
pcs: Set[RTCPeerConnection] = set()      # keep refs so they don’t GC early

class Offer(BaseModel):
    sdp: str
    type: str


# Path to persist stream settings
STREAM_SETTINGS_FILE = os.path.join(os.path.dirname(__file__), 'stream_settings.json')

def load_stream_settings() -> dict:
    """Load persisted stream settings if available."""
    if os.path.exists(STREAM_SETTINGS_FILE):
        try:
            with open(STREAM_SETTINGS_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception:
            return {}
    return {}

def save_stream_settings(settings: dict):
    """Save stream settings to disk."""
    try:
        with open(STREAM_SETTINGS_FILE, 'w', encoding='utf-8') as f:
            json.dump(settings, f, indent=2)
    except Exception as e:
        log.error('Unable to persist stream settings: %s', e)


app = FastAPI()

# Allow CORS for frontend on localhost
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # You can restrict this to ["http://localhost:8501"] if you want
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize global configuration
global_config_group = auracast_config.AuracastConfigGroup()

# Create multicast controller
multicaster: multicast_control.Multicaster | None = None

@app.post("/init")
async def initialize(conf: auracast_config.AuracastConfigGroup):
    """Initializes the broadcasters."""
    global global_config_group
    global multicaster
    try:
        if conf.transport == 'auto':
            serial_devices = glob.glob('/dev/serial/by-id/*')
            log.info('Found serial devices: %s', serial_devices)
            for device in serial_devices:
                if 'usb-ZEPHYR_Zephyr_HCI_UART_sample' in device:
                    log.info('Using: %s', device)
                    conf.transport = f'serial:{device},115200,rtscts'
                    break

        # check again if transport is still auto
        if conf.transport == 'auto':
            HTTPException(status_code=500, detail='No suitable transport found.')

        # initialize the streams dict
        # persist stream settings for later retrieval
        # Derive audio_mode from first BIG audio_source
        first_source = conf.bigs[0].audio_source if conf.bigs else ''
        if first_source.startswith('device:'):
            audio_mode_persist = 'USB'

            input_device = first_source.split(':', 1)[1] if ':' in first_source else 'default'
        elif first_source == 'webrtc':
            audio_mode_persist = 'Webapp'
            input_device = None
        else:
            audio_mode_persist = 'Network'
            input_device = None
        save_stream_settings({
            'channel_names': [big.name for big in conf.bigs],
            'languages': [big.language for big in conf.bigs],
            'audio_mode': audio_mode_persist,
            'input_device': input_device,
            'timestamp': datetime.utcnow().isoformat()
        })
        global_config_group = conf
        # If there is an existing multicaster, cleanly shut it down first so audio devices are released
        if multicaster is not None:
            try:
                await multicaster.shutdown()
            except Exception:
                log.warning("Failed to shutdown previous multicaster", exc_info=True)
            
        log.info(
            'Initializing multicaster with config:\n %s', conf.model_dump_json(indent=2)
        )
        multicaster = multicast_control.Multicaster(
            conf,
            conf.bigs,
        )
        await multicaster.init_broadcast()

        # Auto-start streaming only when using a local USB audio device. For Webapp mode the
        # streamer is started by the /offer handler once the WebRTC track arrives so we know
        # the peer connection is established.
        if any(big.audio_source.startswith("device:") for big in conf.bigs):
            await multicaster.start_streaming()
    except Exception as e:
        log.error("Exception in /init: %s", traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/stream_lc3")
async def send_audio(audio_data: dict[str, str]):
    """Sends a block of pre-coded LC3 audio."""
    if multicaster is None:
        raise HTTPException(status_code=500, detail='Auracast endpoint was never intialized')
    try:
        for big in global_config_group.bigs:    
            assert big.language in audio_data, HTTPException(status_code=500, detail='language len missmatch')
            log.info('Received a send audio request for %s', big.language)
            big.audio_source = audio_data[big.language].encode('latin-1') # TODO: use base64 encoding

        multicaster.big_conf = global_config_group.bigs
        await multicaster.start_streaming()
        return {"status": "audio_sent"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/stop_audio")
async def stop_audio():
    """Stops streaming."""
    try:
        # First close any active WebRTC peer connections so their track loops finish cleanly
        close_tasks = [pc.close() for pc in list(pcs)]
        pcs.clear()
        if close_tasks:
            await asyncio.gather(*close_tasks, return_exceptions=True)

        # Now shut down the multicaster and release audio devices
        running=False
        if multicaster is not None:
            await multicaster.stop_streaming()
            running=True

        return {"status": "stopped", "was_running": running}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/status")
async def get_status():
    """Gets the current status of the multicaster together with persisted stream info."""
    status = multicaster.get_status() if multicaster else {
        'is_initialized': False,
        'is_streaming': False,
    }
    status.update(load_stream_settings())
    return status

@app.get("/audio_inputs")
async def list_audio_inputs():
    """Return available hardware audio input devices for USB mode."""
    try:
        # Re-scan devices on Linux, see https://github.com/spatialaudio/python-sounddevice/issues/16
        if sys.platform == 'linux':
            log.info("Re-initializing sounddevice to scan for new devices")
            sd._terminate()
            sd._initialize()
        devs = sd.query_devices()
        inputs = [
            {"id": idx, "name": d["name"]}
            for idx, d in enumerate(devs)
            if d.get("max_input_channels", 0) > 0 and ("(hw:" in d["name"].lower() or "usb" in d["name"].lower())
        ]
        log.info('Found %d audio input devices:', len(inputs))
        for i in inputs:
            log.info('  %s', i)
        return {"inputs": inputs}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/offer")
async def offer(offer: Offer):
    log.info("/offer endpoint called")

    # If a previous PeerConnection is still alive, close it so we only ever keep one active.
    if pcs:
        log.info("Closing %d existing PeerConnection(s) before creating a new one", len(pcs))
        close_tasks = [p.close() for p in list(pcs)]
        await asyncio.gather(*close_tasks, return_exceptions=True)
        pcs.clear()

    pc = RTCPeerConnection()  # No STUN needed for localhost
    pcs.add(pc)
    id_ = uuid.uuid4().hex[:8]
    log.info(f"{id_}: new PeerConnection")

    # create directory for records - only for testing
    os.makedirs("./records", exist_ok=True)

    # Do NOT start the streamer yet – we'll start it lazily once we actually
    # receive the first audio frame, ensuring WebRTCAudioInput is ready and
    # avoiding race-conditions on restarts.
    @pc.on("track")
    async def on_track(track: MediaStreamTrack):
        log.info(f"{id_}: track {track.kind} received")
        try:
            first = True
            while True:
                frame: av.audio.frame.AudioFrame = await track.recv()          # RTP audio frame (already decrypted)
                if first:
                    log.info(f"{id_}: frame layout={frame.layout}")
                    log.info(f"{id_}: frame format={frame.format}")
                    log.info(
                        f"{id_}: frame sample_rate={frame.sample_rate}, samples_per_channel={frame.samples}, planes={frame.planes}"
                    )
                    # Lazily start the streamer now that we know a track exists.
                    if multicaster.streamer is None:
                        await multicaster.start_streaming()
                        # Yield control so the Streamer coroutine has a chance to
                        # create the WebRTCAudioInput before we push samples.
                        await asyncio.sleep(0)
                    first = False
                    # in stereo case this is interleaved data format
                    frame_array = frame.to_ndarray()
                    log.info(f"array.shape{frame_array.shape}")
                    log.info(f"array.dtype{frame_array.dtype}")
                    log.info(f"frame.to_ndarray(){frame_array}")

                    samples = frame_array.reshape(-1)
                    log.info(f"samples.shape: {samples.shape}")

                    if frame.layout.name == 'stereo':
                        # Interleaved stereo: [L0, R0, L1, R1, ...]
                        mono_array = samples[::2]  # Take left channel
                    else:
                        mono_array = samples

                    log.info(f"mono_array.shape: {mono_array.shape}")


                frame_array = frame.to_ndarray()

                # Flatten in case it's (1, N) or (N,)
                samples = frame_array.reshape(-1)

                if frame.layout.name == 'stereo':
                    # Interleaved stereo: [L0, R0, L1, R1, ...]
                    mono_array = samples[::2]  # Take left channel
                else:
                    mono_array = samples

                # Get current WebRTC audio input (streamer may have been restarted)
                big0 = list(multicaster.bigs.values())[0]
                audio_input = big0.get('audio_input')
                # Wait until the streamer has instantiated the WebRTCAudioInput
                if audio_input is None or getattr(audio_input, 'closed', False):
                    continue
                # Feed mono PCM samples to the global WebRTC audio input
                await audio_input.put_samples(mono_array.astype(np.int16))

                # Save to WAV file - only for testing
                # if not hasattr(pc, 'wav_writer'):
                #     import wave
                #     wav_path = f"./records/auracast_{id_}.wav"
                #     pc.wav_writer = wave.open(wav_path, "wb")
                #     pc.wav_writer.setnchannels(1)  # mono
                #     pc.wav_writer.setsampwidth(2)  # 16-bit PCM
                #     pc.wav_writer.setframerate(frame.sample_rate)

                # pcm_data = mono_array.astype(np.int16).tobytes()
                # pc.wav_writer.writeframes(pcm_data)


        except Exception as e:
            log.error(f"{id_}: Exception in on_track: {e}")
        finally:
            # Always close the wav file when the track ends or on error
            if hasattr(pc, 'wav_writer'):
                try:
                    pc.wav_writer.close()
                except Exception:
                    pass
                del pc.wav_writer

    # --- SDP negotiation ---
    log.info(f"{id_}: setting remote description")
    await pc.setRemoteDescription(RTCSessionDescription(**offer.model_dump()))

    log.info(f"{id_}: creating answer")
    answer = await pc.createAnswer()
    sdp = answer.sdp
    # Insert a=ptime using the global PTIME variable
    ptime_line = f"a=ptime:{PTIME}"
    if "a=sendrecv" in sdp:
        sdp = sdp.replace("a=sendrecv", f"a=sendrecv\n{ptime_line}")
    else:
        sdp += f"\n{ptime_line}"
    new_answer = RTCSessionDescription(sdp=sdp, type=answer.type)
    await pc.setLocalDescription(new_answer)
    log.info(f"{id_}: sending answer with {ptime_line}")
    return {"sdp": pc.localDescription.sdp,
            "type": pc.localDescription.type}


@app.post("/shutdown")
async def shutdown():
    """Stops broadcasting and releases all audio/Bluetooth resources."""
    try:
        await multicaster.shutdown()
        return {"status": "stopped"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == '__main__':
    import uvicorn
    log.basicConfig(
        level=log.DEBUG,
        format='%(module)s.py:%(lineno)d %(levelname)s: %(message)s'
    )
    uvicorn.run(app, host="0.0.0.0", port=5000)