import asyncio
import os
import re
import time
import uuid
import shutil
import aiohttp
from fastapi import FastAPI, BackgroundTasks, Request, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
import shutil
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# tasks = { "task_id": { ... } }
tasks = {}

DATA_DIR = "/data"
os.makedirs(DATA_DIR, exist_ok=True)

class DownloadRequest(BaseModel):
    url: str

async def cleanup_old_files():
    while True:
        try:
            now = time.time()
            for filename in os.listdir(DATA_DIR):
                filepath = os.path.join(DATA_DIR, filename)
                if os.path.isfile(filepath):
                    # 24 hours = 86400 seconds
                    if now - os.path.getmtime(filepath) > 86400:
                        os.remove(filepath)
                        # Remove from tasks if exists
                        for tid, tinfo in list(tasks.items()):
                            if tinfo.get("file_path") == filepath:
                                del tasks[tid]
        except Exception as e:
            print(f"Cleanup error: {e}")
        await asyncio.sleep(3600) # Check every hour

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(cleanup_old_files())

async def download_file(task_id: str, url: str):
    file_path = tasks[task_id]["file_path"]
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url) as response:
                response.raise_for_status()
                total_size = int(response.headers.get('Content-Length', 0))
                tasks[task_id]["total_size"] = total_size
                
                # Check for Content-Disposition header for filename
                cd = response.headers.get('Content-Disposition')
                if cd and 'filename=' in cd:
                    # Extract filename from header
                    fname = re.findall('filename="([^"]+)"', cd)
                    if not fname:
                        fname = re.findall('filename=([^;]+)', cd)
                    if fname:
                        new_filename = fname[0]
                        new_file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
                        tasks[task_id]["file_path"] = new_file_path
                        tasks[task_id]["original_filename"] = new_filename
                        file_path = new_file_path

                downloaded = 0
                start_time = time.time()
                last_time = start_time
                last_downloaded = 0
                
                with open(file_path, 'wb') as f:
                    async for chunk in response.content.iter_chunked(1024 * 1024):
                        if not chunk:
                            break
                        f.write(chunk)
                        downloaded += len(chunk)
                        tasks[task_id]["downloaded"] = downloaded
                        
                        current_time = time.time()
                        if current_time - last_time >= 1.0:
                            speed = (downloaded - last_downloaded) / (current_time - last_time)
                            tasks[task_id]["speed"] = speed
                            last_time = current_time
                            last_downloaded = downloaded
                            
                tasks[task_id]["status"] = "completed"
                tasks[task_id]["speed"] = 0.0
    except Exception as e:
        tasks[task_id]["status"] = "error"
        tasks[task_id]["error"] = str(e)

@app.post("/start_download")
async def start_download(req: DownloadRequest, background_tasks: BackgroundTasks):
    task_id = str(uuid.uuid4())
    
    # Try to extract filename from URL
    filename = req.url.split("/")[-1]
    if "?" in filename:
        filename = filename.split("?")[0]
    if not filename or "." not in filename:
        filename = "downloaded_file.bin"
        
    tasks[task_id] = {
        "task_id": task_id,
        "url": req.url,
        "status": "downloading",
        "total_size": 0,
        "downloaded": 0,
        "speed": 0.0,
        "file_path": os.path.join(DATA_DIR, f"{task_id}_{filename}"),
        "original_filename": filename,
        "timestamp": time.time()
    }
    
    background_tasks.add_task(download_file, task_id, req.url)
    return {"task_id": task_id}

@app.get("/status/{task_id}")
async def get_status(task_id: str):
    if task_id not in tasks:
        return {"error": "Task not found"}
    return tasks[task_id]

@app.get("/download/{task_id}")
async def download(task_id: str):
    if task_id not in tasks or tasks[task_id]["status"] != "completed":
        return {"error": "File not ready"}
    
    file_path = tasks[task_id]["file_path"]
    filename = tasks[task_id].get("original_filename", f"downloaded_{task_id}.bin")
    
    return FileResponse(file_path, filename=filename)

@app.get("/stream/{task_id}")
async def stream(task_id: str, request: Request):
    if task_id not in tasks or tasks[task_id]["status"] != "completed":
        return {"error": "File not ready"}
    
    file_path = tasks[task_id]["file_path"]
    file_size = os.path.getsize(file_path)
    
    range_header = request.headers.get("Range")
    if range_header:
        byte1, byte2 = 0, None
        match = range_header.replace("bytes=", "").split("-")
        if match[0]:
            byte1 = int(match[0])
        if len(match) > 1 and match[1]:
            byte2 = int(match[1])
        
        length = file_size - byte1
        if byte2 is not None:
            length = byte2 + 1 - byte1
            
        def file_iterator(start, length):
            with open(file_path, "rb") as f:
                f.seek(start)
                chunk_size = 1024 * 1024
                while length > 0:
                    read_size = min(chunk_size, length)
                    data = f.read(read_size)
                    if not data:
                        break
                    yield data
                    length -= len(data)
                    
        headers = {
            "Content-Range": f"bytes {byte1}-{byte1+length-1}/{file_size}",
            "Accept-Ranges": "bytes",
            "Content-Length": str(length),
        }
        return StreamingResponse(file_iterator(byte1, length), status_code=206, headers=headers)
    else:
        return FileResponse(file_path)

@app.get("/storage")
async def get_storage():
    total, used, free = shutil.disk_usage(DATA_DIR)
    return {"total": total, "used": used, "free": free}

@app.post("/delete_all")
async def delete_all():
    for filename in os.listdir(DATA_DIR):
        filepath = os.path.join(DATA_DIR, filename)
        try:
            if os.path.isfile(filepath):
                os.remove(filepath)
        except Exception as e:
            print(f"Failed to delete {filepath}: {e}")
    tasks.clear()
    return {"status": "success"}

@app.get("/history")
async def get_history():
    return {"history": list(tasks.values())}


@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    task_id = str(uuid.uuid4())
    original_filename = file.filename if file.filename else "uploaded_file.bin"
    file_path = os.path.join(DATA_DIR, f"{task_id}_{original_filename}")
    
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
        
    file_size = os.path.getsize(file_path)
    
    tasks[task_id] = {
        "task_id": task_id,
        "url": "local_upload",
        "status": "completed",
        "total_size": file_size,
        "downloaded": file_size,
        "speed": 0.0,
        "file_path": file_path,
        "original_filename": original_filename,
        "timestamp": time.time()
    }
    return {"task_id": task_id}


@app.post("/start_gofile_transfer")
async def start_gofile_transfer(req: DownloadRequest, background_tasks: BackgroundTasks):
    task_id = str(uuid.uuid4())
    filename = req.url.split("/")[-1]
    if "?" in filename:
        filename = filename.split("?")[0]
    if not filename or "." not in filename:
        filename = "transfer_file.bin"
        
    tasks[task_id] = {
        "task_id": task_id,
        "url": req.url,
        "status": "downloading",
        "total_size": 0,
        "downloaded": 0,
        "speed": 0.0,
        "file_path": os.path.join(DATA_DIR, f"{task_id}_{filename}"),
        "original_filename": filename,
        "gofile_url": None,
        "timestamp": time.time()
    }
    background_tasks.add_task(process_gofile_transfer, task_id, req.url)
    return {"task_id": task_id}

async def process_gofile_transfer(task_id: str, url: str):
    file_path = tasks[task_id]["file_path"]
    try:
        # 1. Download
        async with aiohttp.ClientSession() as session:
            async with session.get(url) as response:
                response.raise_for_status()
                total_size = int(response.headers.get('Content-Length', 0))
                tasks[task_id]["total_size"] = total_size
                
                cd = response.headers.get('Content-Disposition')
                if cd and 'filename=' in cd:
                    fname = re.findall('filename="([^"]+)"', cd)
                    if not fname:
                        fname = re.findall('filename=([^;]+)', cd)
                    if fname:
                        new_filename = fname[0]
                        new_file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
                        tasks[task_id]["file_path"] = new_file_path
                        tasks[task_id]["original_filename"] = new_filename
                        file_path = new_file_path

                downloaded = 0
                start_time = time.time()
                last_time = start_time
                last_downloaded = 0
                
                with open(file_path, 'wb') as f:
                    async for chunk in response.content.iter_chunked(1024 * 1024):
                        if not chunk:
                            break
                        f.write(chunk)
                        downloaded += len(chunk)
                        tasks[task_id]["downloaded"] = downloaded
                        
                        current_time = time.time()
                        if current_time - last_time >= 1.0:
                            tasks[task_id]["speed"] = (downloaded - last_downloaded) / (current_time - last_time)
                            last_time = current_time
                            last_downloaded = downloaded
        
        # 2. Upload to GoFile
        tasks[task_id]["status"] = "uploading_to_gofile"
        tasks[task_id]["speed"] = 0.0
        
        async with aiohttp.ClientSession() as session:
            # Get server
            async with session.get("https://api.gofile.io/servers") as resp:
                servers_data = await resp.json()
                server = servers_data["data"]["servers"][0]["name"]
            
            # Upload
            upload_url = f"https://{server}.gofile.io/uploadFile"
            with open(file_path, 'rb') as f:
                data = aiohttp.FormData()
                data.add_field('file', f, filename=tasks[task_id]["original_filename"])
                async with session.post(upload_url, data=data) as upload_resp:
                    upload_result = await upload_resp.json()
                    if upload_result["status"] == "ok":
                        tasks[task_id]["gofile_url"] = upload_result["data"]["downloadPage"]
                        tasks[task_id]["status"] = "completed"
                    else:
                        raise Exception("GoFile upload failed")
                        
        # Cleanup local file after successful transfer
        try:
            os.remove(file_path)
        except:
            pass
            
    except Exception as e:
        tasks[task_id]["status"] = "error"
        tasks[task_id]["error"] = str(e)

if __name__ == "__main__":

    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)