package router

import (
	"fmt"
	"sort"
	"sync"
)

type ModelCapability struct {
	ModelID         string
	Provider        string
	MediaType       string
	QualityScore    float64
	MotionHandling  float64
	Composition     float64
	LatencyP50Ms    float64
	LatencyP99Ms    float64
	TokenCostPerRequest float64
}

type UserPreference struct {
	QualityWeight float64
	SpeedWeight   float64
	CostWeight    float64
	RegionRestrict string
	MaxLatencyMs  float64
}

type RouterEngine struct {
	mu          sync.RWMutex
	models      map[string]*ModelCapability
	preferences *UserPreference
}

func NewRouterEngine(preferences *UserPreference) *RouterEngine {
	return &RouterEngine{
		models:      make(map[string]*ModelCapability),
		preferences: preferences,
	}
}

func (e *RouterEngine) RegisterModel(model *ModelCapability) {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.models[model.ModelID] = model
}

func (e *RouterEngine) RouteRequest(reqType string) (*ModelCapability, float64, error) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	
	candidates := make([]*ModelCapability, 0)
	for _, m := range e.models {
		if m.MediaType != reqType { continue }
		if e.preferences.RegionRestrict == "US_only" && 
		   (m.Provider == "bytedance" || m.Provider == "alibaba") { continue }
		if m.LatencyP99Ms > e.preferences.MaxLatencyMs { continue }
		candidates = append(candidates, m)
	}
	
	if len(candidates) == 0 {
		return nil, 0, fmt.Errorf("no suitable model")
	}
	
	type scoredModel struct {
		model *ModelCapability
		score float64
	}
	
	scored := make([]scoredModel, len(candidates))
	for i, m := range candidates {
		qualityScore := m.QualityScore / 100.0
		speedScore := 1.0 - (m.LatencyP50Ms / 10000.0)
		if speedScore < 0 { speedScore = 0 }
		costScore := 1.0 - (m.TokenCostPerRequest / 10.0)
		if costScore < 0 { costScore = 0 }
		
		totalScore := e.preferences.QualityWeight*qualityScore +
			e.preferences.SpeedWeight*speedScore +
			e.preferences.CostWeight*costScore
		scored[i] = scoredModel{model: m, score: totalScore}
	}
	
	sort.Slice(scored, func(i, j int) bool { return scored[i].score > scored[j].score })
	return scored[0].model, scored[0].score, nil
}

"""
Media Quality Evaluation System
"""
class MediaQualityEvaluator:
    def __init__(self):
        self.metrics = {
            "video": {
                "motion_fluidity": 0.25,
                "temporal_consistency": 0.20,
                "scene_composition": 0.15,
                "lighting_quality": 0.10,
                "texture_detail": 0.10,
                "camera_movement": 0.10,
                "color_accuracy": 0.10,
            },
            "image": {
                "composition": 0.25, "sharpness": 0.20,
                "color_harmony": 0.15, "lighting": 0.15,
                "detail_preservation": 0.15, "style_fidelity": 0.10,
            },
        }
    
    def evaluate_model(self, model_name, media_type, test_samples):
        metrics = self.metrics.get(media_type, {})
        scores = {}
        for metric, weight in metrics.items():
            base_score = self._simulate_metric_score(model_name, metric)
            scores[metric] = base_score
        
        weighted = sum(scores[m] * w for m, w in metrics.items())
        scores["overall"] = round(weighted, 1)
        return scores
    
    def _simulate_metric_score(self, model_name, metric):
        profiles = {
            "runway_gen4.5": {"motion_fluidity": 88, "composition": 88},
            "google_veo3": {"motion_fluidity": 92, "composition": 85},
            "bytedance_jimeng2.0": {"motion_fluidity": 85, "composition": 82},
        }
        return profiles.get(model_name, {}).get(metric, 75.0) / 100.0


class PreferenceRouter:
    """Route based on quality, speed, cost, and regional compliance"""
    
    def __init__(self):
        self.latency = {
            "runway_gen4.5": 12000, "google_veo3": 8000,
            "bytedance_jimeng2.0": 10000, "alibaba_wan2.1": 9000,
        }
        self.cost = {
            "runway_gen4.5": 0.35, "google_veo3": 0.45,
            "bytedance_jimeng2.0": 0.08, "alibaba_wan2.1": 0.06,
        }
        self.evaluator = MediaQualityEvaluator()
    
    def route(self, media_type, preferences, region="any"):
        models = ["runway_gen4.5", "google_veo3", "bytedance_jimeng2.0", "alibaba_wan2.1"]
        
        if region == "US_only":
            models = [m for m in models if m.startswith(("runway", "google"))]
        
        best_score, best_model = -1, ""
        for model in models:
            quality = self.evaluator.evaluate_model(model, media_type, [])["overall"]
            speed = max(0, 1.0 - self.latency[model] / 30000.0)
            cost_score = max(0, 1.0 - self.cost.get(model, 1.0) / 1.0)
            
            total = (preferences.get("quality", 0.33) * quality +
                     preferences.get("speed", 0.33) * speed +
                     preferences.get("cost", 0.34) * cost_score)
            if total > best_score:
                best_score, best_model = total, model
        return best_model, round(best_score, 3)


router = PreferenceRouter()
print("Quality-first:", router.route("video", {"quality": 0.7, "speed": 0.2, "cost": 0.1}))
print("Speed-first:", router.route("video", {"quality": 0.2, "speed": 0.6, "cost": 0.2}))
print("Cost-first:", router.route("video", {"quality": 0.2, "speed": 0.2, "cost": 0.6}))
print("US-only:", router.route("video", {"quality": 0.4, "speed": 0.3, "cost": 0.3}, "US_only"))