package extractor

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

type File struct {
	Path    string
	Content string
}

// Extract 从模型响应中提取代码文件
// 使用纯字符串解析，不依赖 goldmark AST
func Extract(response string) ([]File, error) {
	var files []File
	seen := make(map[string]bool)

	// 按 // File: 分割
	parts := strings.Split(response, "// File:")

	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}

		// 第一行是文件名
		lines := strings.SplitN(part, "\n", 2)
		if len(lines) < 2 {
			continue
		}

		fileName := strings.TrimSpace(lines[0])
		code := strings.TrimSpace(lines[1])

		if fileName == "" || code == "" {
			continue
		}

		// 清理代码块标记（移除 ```go 或 ```）
		code = strings.TrimPrefix(code, "```go")
		code = strings.TrimPrefix(code, "```")
		code = strings.TrimSuffix(code, "```")
		code = strings.TrimSpace(code)

		// 智能识别文件类型
		fileName = smartFileName(fileName, code)

		// 规范化路径
		fullPath := normalizePath(fileName)

		// 去重
		if seen[fullPath] {
			continue
		}
		seen[fullPath] = true

		files = append(files, File{
			Path:    fullPath,
			Content: code,
		})
	}

	if len(files) == 0 {
		return nil, fmt.Errorf("未找到任何代码块，请确保响应格式为：// File: 文件名\n```go\n代码\n```")
	}

	return files, nil
}

// smartFileName 智能识别文件名
func smartFileName(fileName, code string) string {
	fileName = strings.TrimSpace(fileName)
	fileName = strings.TrimLeft(fileName, "0123456789.- ")

	// 如果文件名是 code.go 或类似通用名，根据内容推断
	base := filepath.Base(fileName)
	if base == "code.go" || base == "code" || base == "main.go" {
		// 根据内容推断
		if strings.Contains(code, "package sdk_test") {
			return "example_test.go"
		}
		if strings.Contains(code, "module ") {
			return "go.mod"
		}
		if strings.Contains(code, "package sdk") && strings.Contains(code, "func (c *Client)") {
			// 检查是哪个 API 文件
			if strings.Contains(code, "CreateKeyOrder") || strings.Contains(code, "OrderKey") {
				return "api_key.go"
			}
			if strings.Contains(code, "BatchOrder") || strings.Contains(code, "BatchQuery") {
				return "api_batch.go"
			}
			if strings.Contains(code, "QueryKey") || strings.Contains(code, "DiscardKey") {
				return "api_key.go"
			}
		}
		if strings.Contains(code, "package sdk") && strings.Contains(code, "type Client struct") {
			return "client.go"
		}
		if strings.Contains(code, "package sdk") && (strings.Contains(code, "type APIError") || strings.Contains(code, "CodeSuccess")) {
			return "errors.go"
		}
		if strings.Contains(code, "package sdk") && (strings.Contains(code, "ParseRSAPrivateKey") || strings.Contains(code, "AESEncryptECB")) {
			return "crypto.go"
		}
		if strings.Contains(code, "package sdk") && (strings.Contains(code, "type RequestEnvelope") || strings.Contains(code, "type KeyOrderInfo")) {
			return "types.go"
		}
	}

	return fileName
}

func normalizePath(path string) string {
	path = strings.TrimSpace(path)
	path = strings.TrimLeft(path, "0123456789.- ")

	// 确保文件扩展名正确
	if !strings.HasSuffix(path, ".go") && !strings.HasSuffix(path, ".mod") {
		// 如果是 go.mod 或 go.work，保持原样
		if path != "go.mod" && path != "go.work" {
			path = path + ".go"
		}
	}

	// 确保路径以 sdk/ 开头
	//if !strings.HasPrefix(path, "sdk/") && !strings.HasPrefix(path, "output/") {
	//	path = filepath.Join("sdk", path)
	//}
	return path
}

// WriteFiles 将文件写入本地
func WriteFiles(files []File, outputDir string) error {
	for _, f := range files {
		fullPath := filepath.Join(outputDir, f.Path)

		if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil {
			return fmt.Errorf("创建目录失败: %w", err)
		}

		if err := os.WriteFile(fullPath, []byte(f.Content), 0644); err != nil {
			return fmt.Errorf("写入文件失败: %w", fullPath, err)
		}
	}
	return nil
}