import httpx
import ffmpeg
import os
import sqlite3
import uvicorn
import uuid
from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from starlette.responses import FileResponse, JSONResponse
from typing import Optional

# --- Configuration ---
app = FastAPI()
DOWNLOAD_DIR = "downloads"
ENCODED_DIR = "encoded"
DB_FILE = "jobs.db"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)
os.makedirs(ENCODED_DIR, exist_ok=True)
# ---------------------

# --- Database Setup ---
def get_db():
    conn = sqlite3.connect(DB_FILE)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    with get_db() as conn:
        conn.execute('''
            CREATE TABLE IF NOT EXISTS jobs (
                task_id TEXT PRIMARY KEY,
                status TEXT NOT NULL,
                file_name TEXT,
                error_message TEXT,
                duration INTEGER 
            )
        ''')
        conn.commit()

# --- Background Re-encoding Task ---
def process_video_task(url: str, cookie: Optional[str], task_id: str):
    db = get_db()
    temp_in_path = os.path.join(DOWNLOAD_DIR, f"{task_id}_in.mp4")
    temp_out_path = os.path.join(ENCODED_DIR, f"{task_id}_out.mp4")

    try:
        # 1. Update DB: Downloading
        db.execute("UPDATE jobs SET status = 'downloading' WHERE task_id = ?", (task_id,))
        db.commit()
        
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36",
            "Referer": "https://terabox.com/"
        }
        if cookie:
            headers["Cookie"] = cookie 
        
        with httpx.stream("GET", url, headers=headers, follow_redirects=True, timeout=600.0) as r:
            r.raise_for_status()
            with open(temp_in_path, 'wb') as f:
                for chunk in r.iter_bytes(chunk_size=8192):
                    f.write(chunk)
        
        # 2. Update DB: "Fixing" (not encoding)
        db.execute("UPDATE jobs SET status = 'encoding' WHERE task_id = ?", (task_id,))
        db.commit()

        #
        # ✅✅✅ THE FIX IS HERE ✅✅✅
        # We changed this from vcodec='libx264' to 'c=copy'
        # This is the FAST (0% CPU) "re-package"
        #
        (
            ffmpeg
            .input(temp_in_path)
            .output(temp_out_path, c='copy', movflags='+faststart') # <-- CHANGED
            .run(capture_stdout=True, capture_stderr=True)
        )
        
        # 3. Get duration using ffprobe
        probe = ffmpeg.probe(temp_out_path)
        duration = int(float(probe['format']['duration'])) # Get duration in seconds
        
        # 4. Update DB: Complete
        final_file_name = f"{task_id}_out.mp4"
        db.execute(
            "UPDATE jobs SET status = 'complete', file_name = ?, duration = ? WHERE task_id = ?", 
            (final_file_name, duration, task_id)
        )
        db.commit()

    except Exception as e:
        error_msg = str(e)
        if hasattr(e, 'stderr'): 
            error_msg = e.stderr.decode()
        print(f"Error processing task {task_id}: {error_msg}")
        db.execute("UPDATE jobs SET status = 'error', error_message = ? WHERE task_id = ?", (error_msg, task_id))
        db.commit()
    finally:
        if os.path.exists(temp_in_path):
            os.remove(temp_in_path)
        db.close()

# --- API Endpoints ---
class VideoRequest(BaseModel):
    url: str
    cookie: Optional[str] = None 

@app.post("/process")
async def start_processing(request: VideoRequest, background_tasks: BackgroundTasks):
    """ START the encoding job. """
    task_id = str(uuid.uuid4())
    
    with get_db() as db:
        db.execute("INSERT INTO jobs (task_id, status) VALUES (?, 'queued')", (task_id,))
        db.commit()
        
    background_tasks.add_task(process_video_task, request.url, request.cookie, task_id)
    return {"status": "queued", "task_id": task_id}

@app.get("/status/{task_id}")
async def get_status(task_id: str):
    """ CHECK the status of the encoding job. """
    with get_db() as db:
        job = db.execute("SELECT * FROM jobs WHERE task_id = ?", (task_id,)).fetchone()
    
    if not job:
        return JSONResponse(status_code=404, content={"status": "not_found"})
    
    return dict(job) # This will now include 'duration'

@app.get("/download/{file_name}")
async def download_file(file_name: str):
    """ DOWNLOAD the final re-encoded file. """
    file_path = os.path.join(ENCODED_DIR, file_name)
    if not os.path.exists(file_path):
        return JSONResponse(status_code=404, content={"status": "file_not_found"})
    
    return FileResponse(file_path, media_type='video/mp4', filename=file_name)

@app.get("/")
async def health_check():
    """ Health check endpoint. """
    return {"status": "ok"}


# --- Server Start ---
@app.on_event("startup")
async def startup_event():
    init_db() # Create the database and table on startup
    print("Database initialized.")

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