# app.py
import os
import time
import uuid
import asyncio
import libtorrent as lt
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import magic

app = FastAPI()
active_sessions = {}

def get_largest_file(torrent_info):
    files = torrent_info.files()
    file_sizes = [(i, files.file_size(i)) for i in range(files.num_files())]
    if not file_sizes:
        return None
    return max(file_sizes, key=lambda x: x[1])

async def add_torrent(magnet_link: str, save_path: str):
    ses = lt.session()
    params = {
        'save_path': save_path,
        'storage_mode': lt.storage_mode_t(2),
    }

    handle = lt.add_magnet_uri(ses, magnet_link, params)
    
    # Wait for metadata
    while not handle.has_metadata():
        await asyncio.sleep(1)
    
    torrent_info = handle.get_torrent_info()
    file_info = get_largest_file(torrent_info)
    if not file_info:
        raise Exception("No files found in torrent")
    
    file_index, _ = file_info
    file_path = os.path.join(save_path, torrent_info.files().file_path(file_index))
    
    # Prioritize largest file
    for i in range(torrent_info.num_files()):
        handle.file_priority(i, 7 if i == file_index else 0)
    
    handle.set_sequential_download(True)
    ses.resume()
    
    return {
        "session": ses,
        "handle": handle,
        "file_path": file_path,
        "downloading": True
    }

@app.post("/start_stream")
async def start_stream(magnet_link: str):
    torrent_id = str(uuid.uuid4())
    save_path = f"./downloads/{torrent_id}"
    os.makedirs(save_path, exist_ok=True)
    
    try:
        session_info = await add_torrent(magnet_link, save_path)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
    active_sessions[torrent_id] = session_info
    return {"stream_url": f"/stream/{torrent_id}"}

@app.get("/stream/{torrent_id}")
async def stream_torrent(torrent_id: str):
    if torrent_id not in active_sessions:
        raise HTTPException(status_code=404, detail="Torrent session not found")
    
    session_info = active_sessions[torrent_id]
    file_path = session_info["file_path"]
    
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")
    
    mime = magic.Magic(mime=True)
    mime_type = mime.from_file(file_path)
    
    def file_generator():
        last_position = 0
        while True:
            if not os.path.exists(file_path):
                time.sleep(1)
                continue
            
            with open(file_path, "rb") as f:
                f.seek(last_position)
                data = f.read(1024 * 1024)  # 1MB chunks
                if not data:
                    if session_info["handle"].is_seed():
                        break
                    time.sleep(1)
                    continue
                last_position = f.tell()
                yield data
    
    return StreamingResponse(file_generator(), media_type=mime_type)

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