package text2scene

import (
	"encoding/json"
	"fmt"
	"log"
	"math"
	"sync"
)

// Scene 表示生成的3D场景
type Scene struct {
	ID        string          `json:"id"`
	Prompt    string          `json:"prompt"`
	Objects   []*SceneObject  `json:"objects"`
	Lights    []*Light        `json:"lights"`
	Camera    *Camera         `json:"camera"`
	CreatedAt int64           `json:"created_at"`
}

// SceneObject 场景中的物体
type SceneObject struct {
	ID         int              `json:"id"`
	Label      string           `json:"label"`
	Position   [3]float64       `json:"position"`   // [x, y, z]
	Rotation   [3]float64       `json:"rotation"`   // 欧拉角
	Scale      [3]float64       `json:"scale"`      // 缩放因子
	MeshID     string           `json:"mesh_id"`    // 网格资源ID
	Material   *Material        `json:"material"`   // 材质参数
	BoundingBox BoundingBox     `json:"bounding_box"`
}

// Material 物理材质参数
type Material struct {
	BaseColor     [3]float64 `json:"base_color"`    // RGB基础颜色
	Metallic      float64    `json:"metallic"`      // 金属度 0-1
	Roughness     float64    `json:"roughness"`     // 粗糙度 0-1
	NormalMapID   string     `json:"normal_map_id"` // 法线贴图
	OcclusionMapID string    `json:"occlusion_map_id"`
}

// BoundingBox 轴对齐包围盒
type BoundingBox struct {
	Min [3]float64 `json:"min"`
	Max [3]float64 `json:"max"`
}

// TextParser 自然语言解析器
type TextParser struct {
	// 预训练的词嵌入模型（简化版使用词频统计）
	wordVectors map[string][]float64
	// 空间关系映射表
	spatialRelations map[string]func([3]float64, [3]float64) [3]float64
}

// NewTextParser 创建解析器实例
func NewTextParser() *TextParser {
	return &TextParser{
		wordVectors:      loadWordVectors(),
		spatialRelations: initSpatialRelations(),
	}
}

// ParsePrompt 解析文本提示，生成场景布局
func (tp *TextParser) ParsePrompt(prompt string) (*SceneLayout, error) {
	// 1. 分词与词性标注
	tokens := tokenize(prompt)
	
	// 2. 提取物体名称和空间关系
	var objects []ObjectDescription
	var relations []SpatialRelation
	
	for _, token := range tokens {
		// 检查是否为已知物体类别
		if isObjectCategory(token.Text) {
			objects = append(objects, ObjectDescription{
				Category: token.Text,
				Position: [3]float64{0, 0, 0}, // 默认原点
			})
		}
		
		// 检查空间关系词
		if rel, ok := tp.spatialRelations[token.Text]; ok {
			relations = append(relations, SpatialRelation{
				Relation: rel,
				Source:   len(objects) - 2, // 前一个物体
				Target:   len(objects) - 1, // 当前物体
			})
		}
	}
	
	// 3. 构建空间关系图（SRG）
	srg := NewSpatialRelationGraph(objects, relations)
	
	// 4. 使用图神经网络优化布局
	layout, err := optimizeLayout(srg)
	if err != nil {
		return nil, fmt.Errorf("layout optimization failed: %w", err)
	}
	
	return layout, nil
}

// optimizeLayout 基于物理约束优化布局
func optimizeLayout(srg *SpatialRelationGraph) (*SceneLayout, error) {
	const (
		maxIterations = 100
		learningRate  = 0.01
		gravityWeight = 0.3
		collisionWeight = 0.5
	)
	
	layout := &SceneLayout{
		Objects: srg.Objects,
	}
	
	for i := 0; i < maxIterations; i++ {
		// 计算重力约束梯度
		gravityGrad := computeGravityGradient(layout)
		
		// 计算碰撞约束梯度
		collisionGrad := computeCollisionGradient(layout)
		
		// 更新位置
		for j := range layout.Objects {
			for k := 0; k < 3; k++ {
				layout.Objects[j].Position[k] -= learningRate * (
					gravityWeight*gravityGrad[j][k] +
					collisionWeight*collisionGrad[j][k],
				)
			}
		}
		
		// 检查收敛
		if isConverged(gravityGrad, collisionGrad) {
			break
		}
	}
	
	return layout, nil
}

// NeRFRenderer 神经辐射场渲染器
type NeRFRenderer struct {
	// 粗粒度体素网格
	coarseGrid *VoxelGrid
	// 细粒度体素网格
	fineGrid *VoxelGrid
	// 颜色和密度网络
	colorNet *NeuralNetwork
	densityNet *NeuralNetwork
	// 渲染参数
	config RenderConfig
}

// RenderConfig 渲染配置
type RenderConfig struct {
	ImageWidth    int     `json:"image_width"`
	ImageHeight   int     `json:"image_height"`
	FOV           float64 `json:"fov"`           // 视场角（弧度）
	NearPlane     float64 `json:"near_plane"`
	FarPlane      float64 `json:"far_plane"`
	SamplesPerRay int     `json:"samples_per_ray"` // 每条光线采样点数
	BatchSize     int     `json:"batch_size"`
}

// NewNeRFRenderer 创建渲染器
func NewNeRFRenderer(config RenderConfig) *NeRFRenderer {
	return &NeRFRenderer{
		coarseGrid: NewVoxelGrid(32, 32, 32), // 粗粒度网格
		fineGrid:   NewVoxelGrid(256, 256, 256),
		colorNet:   NewColorNetwork(),
		densityNet: NewDensityNetwork(),
		config:     config,
	}
}

// RenderScene 渲染场景到2D图像
func (r *NeRFRenderer) RenderScene(scene *Scene, camera *Camera) ([][][3]float64, error) {
	// 1. 计算光线方向
	rays := computeCameraRays(camera, r.config)
	
	// 2. 并行渲染（使用goroutine）
	image := make([][][3]float64, r.config.ImageHeight)
	var wg sync.WaitGroup
	
	// 将图像分为若干块并行处理
	blockSize := 32
	numBlocksY := (r.config.ImageHeight + blockSize - 1) / blockSize
	numBlocksX := (r.config.ImageWidth + blockSize - 1) / blockSize
	
	for by := 0; by < numBlocksY; by++ {
		for bx := 0; bx < numBlocksX; bx++ {
			wg.Add(1)
			go func(blockY, blockX int) {
				defer wg.Done()
				r.renderBlock(rays, image, blockY, blockX, blockSize)
			}(by, bx)
		}
	}
	
	wg.Wait()
	return image, nil
}

// renderBlock 渲染图像块
func (r *NeRFRenderer) renderBlock(rays [][]Ray, image [][][3]float64, blockY, blockX, blockSize int) {
	startY := blockY * blockSize
	endY := min(startY+blockSize, r.config.ImageHeight)
	startX := blockX * blockSize
	endX := min(startX+blockSize, r.config.ImageWidth)
	
	for y := startY; y < endY; y++ {
		image[y] = make([][3]float64, r.config.ImageWidth)
		for x := startX; x < endX; x++ {
			// 体渲染积分
			color := r.volumeRender(rays[y][x], scene)
			image[y][x] = color
		}
	}
}

// volumeRender 体渲染核心算法
func (r *NeRFRenderer) volumeRender(ray Ray, scene *Scene) [3]float64 {
	var accumulatedColor [3]float64
	var accumulatedTransmittance float64 = 1.0
	
	// 沿光线采样
	tValues := linspace(ray.Near, ray.Far, r.config.SamplesPerRay)
	
	for _, t := range tValues {
		// 计算采样点位置
		point := ray.Origin.Add(ray.Direction.Scale(t))
		
		// 查询密度和颜色
		density := r.queryDensity(point, scene)
		color := r.queryColor(point, ray.Direction, scene)
		
		// 体渲染积分
		alpha := 1.0 - math.Exp(-density*stepSize)
		accumulatedColor = accumulatedColor.Add(
			color.Scale(alpha * accumulatedTransmittance),
		)
		accumulatedTransmittance *= (1.0 - alpha)
		
		// 早期终止（透明度接近0）
		if accumulatedTransmittance < 0.001 {
			break
		}
	}
	
	return accumulatedColor
}

// queryDensity 查询体积密度（使用分层网格加速）
func (r *NeRFRenderer) queryDensity(point [3]float64, scene *Scene) float64 {
	// 1. 先查询粗粒度网格
	coarseDensity := r.coarseGrid.Sample(point)
	
	// 如果粗粒度密度为0，直接返回
	if coarseDensity < 0.01 {
		return 0
	}
	
	// 2. 查询细粒度网格
	fineDensity := r.fineGrid.Sample(point)
	
	// 3. 通过神经网络细化
	features := r.densityNet.Forward(point, scene)
	return fineDensity * features[0]
}

// PhysicsEngine 物理约束引擎
type PhysicsEngine struct {
	gravity [3]float64 // 重力方向
}

// CheckConstraints 检查场景物理合理性
func (pe *PhysicsEngine) CheckConstraints(scene *Scene) []ConstraintViolation {
	var violations []ConstraintViolation
	
	// 1. 重力约束检查
	for _, obj := range scene.Objects {
		// 检查物体是否悬浮（底部离地超过阈值）
		bottomY := obj.BoundingBox.Min[1]
		if bottomY > 0.01 {
			violations = append(violations, ConstraintViolation{
				ObjectID: obj.ID,
				Type:     GravityViolation,
				Message:  fmt.Sprintf("物体 %s 悬浮在高度 %.2f", obj.Label, bottomY),
			})
		}
	}
	
	// 2. 碰撞检测
	for i := 0; i < len(scene.Objects); i++ {
		for j := i + 1; j < len(scene.Objects); j++ {
			if boxesIntersect(scene.Objects[i].BoundingBox, scene.Objects[j].BoundingBox) {
				violations = append(violations, ConstraintViolation{
					ObjectID: scene.Objects[i].ID,
					Type:     CollisionViolation,
					Message:  fmt.Sprintf("物体 %s 与 %s 发生碰撞", 
						scene.Objects[i].Label, scene.Objects[j].Label),
				})
			}
		}
	}
	
	// 3. 光照一致性检查
	for _, light := range scene.Lights {
		// 检查光源是否在物体内部
		for _, obj := range scene.Objects {
			if pointInBox(light.Position, obj.BoundingBox) {
				violations = append(violations, ConstraintViolation{
					ObjectID: light.ID,
					Type:     LightInsideObject,
					Message:  fmt.Sprintf("光源位于物体 %s 内部", obj.Label),
				})
			}
		}
	}
	
	return violations
}

// boxesIntersect 判断两个AABB是否相交
func boxesIntersect(a, b BoundingBox) bool {
	return a.Min[0] < b.Max[0] && a.Max[0] > b.Min[0] &&
		a.Min[1] < b.Max[1] && a.Max[1] > b.Min[1] &&
		a.Min[2] < b.Max[2] && a.Max[2] > b.Min[2]
}

// SceneGenerator 场景生成器主类
type SceneGenerator struct {
	parser      *TextParser
	renderer    *NeRFRenderer
	physics     *PhysicsEngine
	materialGen *MaterialGenerator
}

// GenerateScene 从文本生成完整场景
func (sg *SceneGenerator) GenerateScene(prompt string) (*Scene, error) {
	log.Printf("开始生成场景，提示词: %s", prompt)
	
	// 1. 解析文本布局
	layout, err := sg.parser.ParsePrompt(prompt)
	if err != nil {
		return nil, fmt.Errorf("布局解析失败: %w", err)
	}
	log.Printf("布局解析完成，包含 %d 个物体", len(layout.Objects))
	
	// 2. 为每个物体生成材质
	for i, obj := range layout.Objects {
		material, err := sg.materialGen.GenerateMaterial(obj.Category, prompt)
		if err != nil {
			log.Printf("物体 %s 材质生成失败: %v", obj.Category, err)
			continue
		}
		layout.Objects[i].Material = material
	}
	
	// 3. 构建场景图
	scene := &Scene{
		ID:      generateUUID(),
		Prompt:  prompt,
		Objects: layout.Objects,
		Lights:  layout.Lights,
	}
	
	// 4. 物理约束检查与修正
	violations := sg.physics.CheckConstraints(scene)
	if len(violations) > 0 {
		log.Printf("发现 %d 个物理约束违规，尝试自动修正", len(violations))
		scene = sg.autoFixConstraints(scene, violations)
	}
	
	// 5. 渲染预览（可选）
	preview, err := sg.renderer.RenderScene(scene, layout.Camera)
	if err != nil {
		log.Printf("预览渲染失败: %v", err)
	}
	scene.PreviewImage = preview
	
	log.Printf("场景生成完成，ID: %s", scene.ID)
	return scene, nil
}

// autoFixConstraints 自动修复物理约束违规
func (sg *SceneGenerator) autoFixConstraints(scene *Scene, violations []ConstraintViolation) *Scene {
	// 实现自动修正逻辑
	// 例如：将悬浮物体下降到地面，调整碰撞物体的位置
	return scene
}