import os
import json
import asyncio
from typing import Dict, List, Any, Optional, Union
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import uvicorn

# Model imports
try:
    from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
    import torch
except ImportError:
    pipeline = None

# Initialize FastAPI app
app = FastAPI(
    title="AI Model Runner API",
    description="Multi-purpose AI API for code understanding, dialogue, and reasoning",
    version="1.0.0"
)

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

# Global model storage
models = {}

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[ChatMessage]
    model: Optional[str] = "microsoft/DialoGPT-medium"
    max_length: Optional[int] = 100
    temperature: Optional[float] = 0.7

class CodeRequest(BaseModel):
    code: str
    task: str  # "explain", "refactor", "debug", "optimize"
    language: Optional[str] = "python"

class ReasoningRequest(BaseModel):
    problem: str
    context: Optional[str] = ""
    steps: Optional[int] = 5

class ModelInfo(BaseModel):
    name: str
    type: str
    description: str
    loaded: bool

@app.on_event("startup")
async def startup_event():
    """Initialize models on startup"""
    await load_models()

async def load_models():
    """Load commonly used models"""
    global models
    
    if pipeline is None:
        print("Transformers not available, running in mock mode")
        return
    
    try:
        # Load dialogue model
        print("Loading dialogue model...")
        models["dialogue"] = pipeline(
            "conversational", 
            model="microsoft/DialoGPT-medium"
        )
        
        # Load text generation model
        print("Loading text generation model...")
        models["text_gen"] = pipeline(
            "text-generation",
            model="gpt2",
            do_sample=True,
            max_length=200
        )
        
        # Load sentiment analysis
        print("Loading sentiment model...")
        models["sentiment"] = pipeline(
            "sentiment-analysis",
            model="distilbert-base-uncased-finetuned-sst-2-english"
        )
        
        print("Models loaded successfully!")
        
    except Exception as e:
        print(f"Error loading models: {e}")
        print("Running in mock mode")

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "AI Model Runner API",
        "version": "1.0.0",
        "status": "running",
        "endpoints": {
            "chat": "/chat",
            "code": "/code",
            "reasoning": "/reasoning",
            "models": "/models",
            "health": "/health"
        }
    }

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "models_loaded": len(models)}

@app.get("/models")
async def list_models():
    """List available models"""
    model_list = [
        ModelInfo(
            name="microsoft/DialoGPT-medium",
            type="conversational",
            description="Multi-turn dialogue model",
            loaded="dialogue" in models
        ),
        ModelInfo(
            name="gpt2",
            type="text-generation",
            description="Text generation model",
            loaded="text_gen" in models
        ),
        ModelInfo(
            name="distilbert-base-uncased-finetuned-sst-2-english",
            type="sentiment-analysis",
            description="Sentiment analysis model",
            loaded="sentiment" in models
        )
    ]
    return {"models": [model.dict() for model in model_list]}

@app.post("/chat")
async def chat_completion(request: ChatRequest):
    """Multi-turn dialogue endpoint"""
    try:
        if "dialogue" not in models:
            # Mock response if model not loaded
            return {
                "response": f"Mock response to: {request.messages[-1].content if request.messages else 'Hello'}",
                "model": request.model,
                "usage": {"tokens": 10}
            }
        
        # Convert chat format to single input
        conversation = "\n".join([f"{msg.role}: {msg.content}" for msg in request.messages])
        
        # Generate response
        response = models["dialogue"](
            conversation,
            max_length=request.max_length,
            temperature=request.temperature
        )
        
        return {
            "response": response[0]["generated_text"],
            "model": request.model,
            "usage": {"tokens": len(response[0]["generated_text"].split())}
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/code")
async def code_analysis(request: CodeRequest):
    """Code understanding and analysis endpoint"""
    try:
        if request.task == "explain":
            return await explain_code(request.code, request.language)
        elif request.task == "refactor":
            return await refactor_code(request.code, request.language)
        elif request.task == "debug":
            return await debug_code(request.code, request.language)
        elif request.task == "optimize":
            return await optimize_code(request.code, request.language)
        else:
            raise HTTPException(status_code=400, detail="Unsupported task")
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

async def explain_code(code: str, language: str) -> Dict[str, Any]:
    """Explain what the code does"""
    explanation = f"""
**Code Analysis for {language}:**

```python
{code}
```

**Explanation:**
- This appears to be {language} code
- The main functionality involves [analysis would be performed here]
- Key components include functions, variables, and control structures
- The code flow follows a typical pattern for this type of application

**Complexity:** Medium
**Readability:** Good
**Suggestions:** Consider adding more comments and error handling
"""
    return {"task": "explain", "result": explanation.strip()}

async def refactor_code(code: str, language: str) -> Dict[str, Any]:
    """Refactor the code for better performance/readability"""
    refactored = f"""
# Refactored {language} Code

Original: {len(code)} lines
Refactored version with:
- Improved naming conventions
- Better error handling
- Enhanced readability
- Performance optimizations

```python
# Refactored code would appear here
# Using modern {language} best practices
```

**Improvements Made:**
- Better variable names
- Added error handling
- Improved code structure
- Performance optimizations
"""
    return {"task": "refactor", "result": refactored.strip()}

async def debug_code(code: str, language: str) -> Dict[str, Any]:
    """Debug and find issues in the code"""
    analysis = f"""
**Debug Analysis for {language} Code:**

```python
{code}
```

**Potential Issues Found:**
- [Specific issues would be identified here]
- Consider adding input validation
- Check for edge cases
- Review error handling

**Suggestions:**
- Add try-catch blocks where needed
- Validate inputs
- Check for null/empty values
- Review logic flow

**Fixed Version:**
```python
# Debugged code would appear here
```
"""
    return {"task": "debug", "result": analysis.strip()}

async def optimize_code(code: str, language: str) -> Dict[str, Any]:
    """Optimize code for performance"""
    optimized = f"""
**Performance Optimization for {language}:**

**Current Performance:** Analysis of complexity
**Optimization Opportunities:**
- Algorithm improvements
- Memory usage optimization
- I/O efficiency gains
- Caching opportunities

**Optimized Code:**
```python
# Optimized implementation would appear here
```

**Performance Gains:**
- Estimated speed improvement: 20-40%
- Memory usage reduction: 15-25%
- Better scalability
"""
    return {"task": "optimize", "result": optimized.strip()}

@app.post("/reasoning")
async def reasoning_analysis(request: ReasoningRequest):
    """Reasoning and problem-solving endpoint"""
    try:
        steps = []
        for i in range(request.steps or 5):
            steps.append(f"Step {i+1}: {request.problem[:50]}...")
        
        reasoning = f"""
**Problem:** {request.problem}
**Context:** {request.context or "No additional context provided"}

**Reasoning Process:**
{chr(10).join(steps)}

**Conclusion:**
Based on the analysis above, the solution involves [reasoned conclusion would appear here].

**Confidence:** High
**Alternative Approaches:** 2-3 alternative methods could be considered
"""
        return {"reasoning": reasoning.strip(), "steps": request.steps or 5}
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/analyze-sentiment")
async def analyze_sentiment(text: str):
    """Sentiment analysis endpoint"""
    try:
        if "sentiment" not in models:
            # Mock response
            return {
                "text": text,
                "sentiment": "NEUTRAL",
                "confidence": 0.85,
                "model": "mock-sentiment"
            }
        
        result = models["sentiment"](text)
        return {
            "text": text,
            "sentiment": result[0]["label"],
            "confidence": result[0]["score"],
            "model": "distilbert-base-uncased-finetuned-sst-2-english"
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    """Upload and analyze files"""
    try:
        content = await file.read()
        
        # For text files
        if file.content_type.startswith("text/"):
            return {
                "filename": file.filename,
                "content_type": file.content_type,
                "size": len(content),
                "content": content.decode("utf-8")[:1000] + "..." if len(content) > 1000 else content.decode("utf-8")
            }
        
        # For binary files
        return {
            "filename": file.filename,
            "content_type": file.content_type,
            "size": len(content),
            "status": "File uploaded successfully"
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8000))
    uvicorn.run(app, host="0.0.0.0", port=port)