from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
import os
import threading
import torch
from app.job_manager import JobManager
from agent.inference import InferenceEngine

app = FastAPI()

# Enable CORS (Critical for cross-port communication)
# This server runs on port 8002
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class TrainRequest(BaseModel):
    data_path: str = None
    manual_text: str = None
    category: str = "text"
    noise_level: float = 0.0

class AutoTrainRequest(BaseModel):
    topic: str = None
    url: str = None
    category: str = "text"

class GrammarCheckRequest(BaseModel):
    text: str

class ParallelTrainRequest(BaseModel):
    topics: list[str]

@app.post("/api/train")
async def trigger_training(request: TrainRequest):
    if not request.data_path and not request.manual_text:
         return {"status": "error", "message": "Either data path or manual text must be provided."}
    
    path = request.data_path if request.data_path else "manual_input"
    if request.data_path and not os.path.exists(path):
         return {"status": "error", "message": "Path not found."}
    
    manager = JobManager()
    job = manager.create_job("custom", path)
    
    def run_train():
        try:
            print(f"Loading data from {path}...")
            print("Starting training...")
            import train.train
            import importlib
            importlib.reload(train.train)
            
            # Pass the job object and category
            train.train.train(
                dataset_path=request.data_path, 
                text_content=request.manual_text,
                job=job, 
                category=request.category,
                noise_level=request.noise_level
            )
            
            print("Training Complete. Model saved to sail.pt")
        except Exception as e:
            print(f"Training Failed: {e}")
            job.status = "failed"
            job.error = str(e)
            job.message = f"Failed: {e}"
            
    thread = threading.Thread(target=run_train)
    job.thread = thread
    thread.start()
    
    return {"status": "success", "message": "Training started.", "job_id": job.id}

@app.post("/api/train/auto")
async def trigger_auto_training(request: AutoTrainRequest):
    if not request.topic and not request.url:
        return {"status": "error", "message": "Either topic or url must be provided."}

    manager = JobManager()
    details = request.topic if request.topic else request.url
    job = manager.create_job("auto", details)

    def run_auto_train():
        try:
            from agent.data_collector import DataCollector
            collector = DataCollector()
            
            text_data = ""
            job.message = "Collecting data..."
            if request.topic:
                print(f"Collecting data for topic: {request.topic}")
                text_data = collector.search_and_collect(request.topic, category=request.category)
            elif request.url:
                print(f"Collecting data from URL: {request.url}")
                text_data = collector.collect_from_url(request.url)
            
            if not text_data:
                print("No data collected.")
                job.status = "failed"
                job.error = "No data collected"
                job.message = "Failed: No data collected"
                return

            print("Starting automatic training...")
            import train.train
            import importlib
            importlib.reload(train.train)
            
            train.train.train(text_content=text_data, job=job, category=request.category)
            print("Auto Training Complete.")
            
        except Exception as e:
            print(f"Auto Training Failed: {e}")
            job.status = "failed"
            job.error = str(e)
            job.message = f"Failed: {e}"

    thread = threading.Thread(target=run_auto_train)
    job.thread = thread
    thread.start()

    return {"status": "success", "message": "Automatic training started.", "job_id": job.id}

@app.post("/api/train/wordnet")
def trigger_wordnet_training(request: AutoTrainRequest):
    """Generates grammar instructions from WordNet and trains."""
    manager = JobManager()
    job = manager.create_job("wordnet", f"Grammar Training (WordNet)")
    
    def run_wordnet_train():
        try:
            from train.wordnet_loader import WordNetLoader
            import json
            
            loader = WordNetLoader()
            job.message = "Loading WordNet (slower first time)..."
            
            # Generate dataset
            print("Generating WordNet dataset...")
            data = loader.generate_instruction_dataset(num_samples=2000) 
            
            if not data:
                job.status = "failed"
                job.error = "No WordNet data generated"
                return

            # Save to JSON
            dataset_path = "wordnet_grammar.json"
            with open(dataset_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2)
                
            print(f"Saved {len(data)} grammar instructions to {dataset_path}")
            
            print("Starting grammar training...")
            import train.train
            import importlib
            importlib.reload(train.train)
            
            train.train.train(dataset_path=dataset_path, job=job, category="grammar")
            print("Grammar Training Complete")
            
        except Exception as e:
            print(f"Grammar Training Failed: {e}")
            job.status = "failed"
            job.error = str(e)
            job.message = f"Failed: {e}"

    thread = threading.Thread(target=run_wordnet_train)
    job.thread = thread
    thread.start()
    
    return {"status": "started", "job_id": job.id, "message": "Grammar training started."}

@app.post("/api/train/instruction")
def trigger_instruction_training(request: AutoTrainRequest):
    """Generates Q&A pairs and trains the model on them."""
    manager = JobManager()
    job = manager.create_job("instruction", f"Category: {request.category}")
    
    def run_instruction_train():
        try:
            from agent.data_collector import DataCollector
            import json
            
            collector = DataCollector()
            job.message = f"Generating {request.category} instructions..."
            print(f"Generating instructions for {request.category}...")
            
            instructions = collector.collect_instructions(request.category)
            
            if not instructions:
                job.status = "failed"
                job.error = "No instructions generated"
                return

            dataset_path = "instructions.json"
            with open(dataset_path, 'w', encoding='utf-8') as f:
                json.dump(instructions, f, indent=2)
                
            print(f"Saved {len(instructions)} instructions to {dataset_path}")
            
            print("Starting instruction training...")
            import train.train
            import importlib
            importlib.reload(train.train)
            
            train.train.train(dataset_path=dataset_path, job=job, category=request.category)
            print("Instruction Training Complete")
            
        except Exception as e:
            print(f"Instruction Training Failed: {e}")
            job.status = "failed"
            job.error = str(e)
            job.message = f"Failed: {e}"

    thread = threading.Thread(target=run_instruction_train)
    job.thread = thread
    thread.start()
    
    return {"status": "started", "job_id": job.id, "message": "Instruction training started."}

@app.post("/api/train/parallel")
def trigger_parallel_training(request: ParallelTrainRequest):
    """Hits the new Multi-Stream single-GPU orchestrator to train multiple topics."""
    if not request.topics or len(request.topics) == 0:
        return {"status": "error", "message": "List of topics required."}
        
    manager = JobManager()
    
    def run_parallel_train():
        from agent.data_collector import DataCollector
        import train.train
        import importlib
        importlib.reload(train.train)
        
        collector = DataCollector()
        tasks = []
        
        for topic in request.topics:
            job = manager.create_job("parallel", f"Parallel: {topic}")
            job.message = f"Collecting {topic}..."
            print(f"Collecting data for {topic}...")
            
            data = collector.search_and_collect(topic, category="custom")
            if data:
                tasks.append({
                    "text_content": data,
                    "category": f"parallel_{topic}",
                    "job": job
                })
            else:
                job.status = "failed"
                job.error = "No data found."
                job.message = "Failed: No Data."
                
        if len(tasks) > 0:
            print(f"Dispatching {len(tasks)} tasks to concurrent CUDA Streams...")
            train.train.train_parallel(tasks)
        else:
            print("No parallel tasks gathered data.")

    thread = threading.Thread(target=run_parallel_train)
    thread.start()
    
    return {"status": "started", "message": f"Parallel multi-stream training started for {len(request.topics)} topics."}

@app.get("/api/jobs")
async def get_jobs():
    manager = JobManager()
    return {"jobs": manager.get_all_jobs()}

@app.post("/api/jobs/{job_id}/{action}")
async def control_job(job_id: str, action: str):
    manager = JobManager()
    success = False
    if action == "pause":
        success = manager.pause_job(job_id)
    elif action == "resume":
        success = manager.resume_job(job_id)
    elif action == "terminate":
        success = manager.terminate_job(job_id)
    
    if success:
        return {"status": "success", "message": f"Job {action}d"}
    else:
        return {"status": "error", "message": f"Failed to {action} job"}

@app.post("/api/reset")
async def reset_brain():
    """Deletes the model file to start fresh."""
    import glob
    try:
        if os.path.exists("sail.pt"):
            os.remove("sail.pt")
        
        for f in glob.glob("sail_*.pt"):
            os.remove(f)
            
        print("Brain reset successfully.")
        return {"status": "success", "message": "Brain reset. Please reload the page."}
    except Exception as e:
        return {"status": "error", "message": str(e)}

@app.get("/api/browse")
async def browse_folder():
    """Opens a native folder picker dialog on the server and returns the path."""
    import subprocess
    try:
        # Run a separate process to avoid Tkinter main thread issues in FastAPI
        cmd = [
            "python", "-c",
            "import tkinter.filedialog as fd; import tkinter; root = tkinter.Tk(); root.withdraw(); root.attributes('-topmost', True); print(fd.askdirectory()); root.destroy()"
        ]
        result = subprocess.run(cmd, capture_output=True, text=True)
        path = result.stdout.strip()
        return {"path": path}
    except Exception as e:
        return {"path": "", "error": str(e)}

@app.get("/api/status")
async def status():
    manager = JobManager()
    active_jobs = [j for j in manager.jobs.values() if j.status == "running"]
    return {
        "is_training": len(active_jobs) > 0,
        "model_exists": os.path.exists("sail.pt"),
        "active_jobs_count": len(active_jobs)
    }