import os
import re
import asyncio
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse
from pyrogram import Client
import uvicorn
from contextlib import asynccontextmanager

# --- Config ---
API_ID = os.getenv("API_ID")
API_HASH = os.getenv("API_HASH")
STRING_SESSION = os.getenv("STRING_SESSION")
ACCESS_KEY = os.getenv("ACCESS_KEY", "0000")

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Bot ကို စတင်ခြင်း
    user_bot = Client(
        "tg_streamer",
        api_id=API_ID,
        api_hash=API_HASH,
        session_string=STRING_SESSION,
        in_memory=True,
        workers=16
    )
    await user_bot.start()
    app.state.tg = user_bot
    print("🚀 Ultimate Smooth Streamer is Online!")
    
    # --- Entity Cache ပြဿနာဖြေရှင်းရန် (Warm-up လုပ်ခြင်း) ---
    print("🔄 Warming up Entity Cache (Fetching Dialogs)...")
    try:
        # အကောင့်ထဲရှိ Group/Channel ၁၀၀ ခုကို အရင်ဖတ်ခိုင်းပြီး မှတ်ထားစေပါသည်
        async for dialog in user_bot.get_dialogs(limit=100):
            pass
        print("✅ Warm-up complete! Entity Cache is ready.")
    except Exception as e:
        print(f"⚠️ Warm-up တွင် အခက်အခဲရှိနေပါသည်: {e}")

    yield
    
    # Bot ကို ပိတ်ခြင်း
    await user_bot.stop()

app = FastAPI(lifespan=lifespan)

# --- Health Check Endpoint ---
@app.get("/")
async def health_check():
    return JSONResponse(content={"status": "running", "message": "Server is alive!"})

def parse_tg_link(link: str):
    match = re.search(r't\.me/(?:c/)?([^/]+)/(\d+)', link)
    if match:
        chat_id = match.group(1)
        if chat_id.isdigit():
            chat_id = int("-100" + chat_id)
        return chat_id, int(match.group(2))
    return None, None

@app.get("/stream")
async def stream_telegram(request: Request, url: str, key: str = None):
    if key != ACCESS_KEY:
        raise HTTPException(status_code=403, detail="Invalid Key")

    chat_id, msg_id = parse_tg_link(url)
    if not chat_id:
        raise HTTPException(status_code=400, detail="Invalid URL")

    client = request.app.state.tg
    
    # --- Streaming မလုပ်ခင် Chat ကို အရင် Fetch လုပ်ခြင်း ---
    try:
        # Chat ID အသစ်ဖြစ်နေပါက Telegram မှ ကြိုတင်အချက်အလက်တောင်းယူခြင်း
        await client.get_chat(chat_id)
    except Exception as e:
        print(f"⚠️ Chat ကို ကြိုတင်ရှာဖွေရာတွင် အခက်အခဲရှိသည် (သို့သော် ဆက်လက်ကြိုးစားမည်): {e}")

    try:
        msg = await client.get_messages(chat_id, msg_id)
        # Media ရှိမရှိ စစ်ဆေးခြင်း
        media = msg.video or msg.document or msg.audio or msg.animation
        
        if not media:
            print(f"❌ Media Not Found in Message: {msg_id}")
            raise Exception("No media found")
            
    except Exception as e:
        print(f"❌ Error fetching message {msg_id}: {str(e)}")
        raise HTTPException(status_code=404, detail=f"File not found or Access Denied: {str(e)}")

    file_size = media.file_size
    range_header = request.headers.get("range")
    
    start_byte = 0
    if range_header:
        match = re.search(r'bytes=(\d+)-', range_header)
        if match:
            start_byte = int(match.group(1))

    async def refined_generator():
        # Telegram ရဲ့ offset က MB နဲ့ သွားတာမို့လို့ပါ
        offset_in_mb = start_byte // (1024 * 1024)
        skip_bytes = start_byte % (1024 * 1024)
        
        try:
            generator = client.stream_media(msg, offset=offset_in_mb, limit=30)
            first_chunk = True
            async for chunk in generator:
                if first_chunk:
                    if len(chunk) > skip_bytes:
                        yield chunk[skip_bytes:]
                    first_chunk = False
                else:
                    yield chunk
        except Exception as e:
            print(f"⚠️ Stream interrupted for {msg_id}: {e}")
            return

    headers = {
        "Content-Type": getattr(media, "mime_type", "video/mp4"),
        "Accept-Ranges": "bytes",
        "Access-Control-Allow-Origin": "*",
        "Cache-Control": "public, max-age=3600",
    }

    if range_header:
        headers["Content-Range"] = f"bytes {start_byte}-{file_size-1}/{file_size}"
        status_code = 206
    else:
        status_code = 200

    return StreamingResponse(
        refined_generator(),
        status_code=status_code,
        headers=headers
    )

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)